Adjust default behavior of generate_chain_data_cache.py

scripts/generate_chain_data_cache.py

@@ -39,7 +39,13 @@ def parse_file(
             local_data["seq"] = seq
             local_data["resolution"] = mmcif.header["resolution"]
            
-            cluster_size = chain_cluster_size_dict.get(full_name.upper(), -1)
+            if(chain_cluster_size_dict is not None):
+                cluster_size = chain_cluster_size_dict.get(
+                    full_name.upper(), -1
+                )
+            else:
+                cluster_size = None
+            
             local_data["cluster_size"] = cluster_size
     elif(ext == ".pdb"):
         with open(os.path.join(args.data_dir, f), "r") as fp:
@@ -53,12 +59,15 @@ def parse_file(
         )
         local_data["resolution"] = 0.
 
-        cluster_size = chain_cluster_size_dict.get(file_id.upper(), None)
-        if(cluster_size is None):
-            print(file_id)
-            return {}
+        cluster_size = chain_cluster_size_dict.get(file_id.upper(), -1)
+        if(chain_cluster_size_dict is not None):
+            cluster_size = chain_cluster_size_dict.get(
+                full_name.upper(), -1
+            )
         else:
-            local_data["cluster_size"] = cluster_size
+            cluster_size = None
+        
+        chain_dict["cluster_size"] = cluster_size
 
         out = {file_id: chain_dict}
 
@@ -66,8 +75,9 @@ def parse_file(
 
 
 def main(args):
-    chain_cluster_size_dict = {}
+    chain_cluster_size_dict = None
     if(args.cluster_file is not None):
+        chain_cluster_size_dict = {}
         with open(args.cluster_file, "r") as fp:
             clusters = [l.strip() for l in fp.readlines()]