Initial fix multipe structures

.gitignore

@@ -1 +1,3 @@
 .DS_Store
+*.pyc
+*.cif

main.py

@@ -40,14 +40,15 @@ def main():
                 if key.lower().startswith("data"):
                     st.markdown(cif_data[key])
             st.markdown(cif_data["_publ_section_comment"])
-            if "_diffrn_ambient_temperature" in cif_data.keys():
-                temperature = float(cif_data["_diffrn_ambient_temperature"].split("(")[0])
-            elif "_cell_measurement_temperature" in cif_data.keys():
-                temperature = float(cif_data["_cell_measurement_temperature"].split("(")[0])
-            else:
-                raise ValueError("Temperature not found in the CIF file. \
-                                    Please provide a temperature in the field _diffrn_ambient_temperature o in the field _cell_measurement_temperature from the CIF file.")
-            st.success("CIF file successfully read.")
+            temperature = 100
+            # if "_diffrn_ambient_temperature" in cif_data.keys():
+            #     temperature = float(cif_data["_diffrn_ambient_temperature"].split("(")[0])
+            # elif "_cell_measurement_temperature" in cif_data.keys():
+            #     temperature = float(cif_data["_cell_measurement_temperature"].split("(")[0])
+            # else:
+            #     raise ValueError("Temperature not found in the CIF file. \
+            #                         Please provide a temperature in the field _diffrn_ambient_temperature o in the field _cell_measurement_temperature from the CIF file.")
+            # st.success("CIF file successfully read.")
 
             data = Data()
             data.x = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int32)

main_local.py

@@ -7,6 +7,7 @@ import torch
 from models.master import create_model
 from process import process_data
 from utils import radius_graph_pbc
+from io import BytesIO, StringIO
 import gc
 
 MEAN_TEMP = torch.tensor(192.1785) #training temp mean
@@ -24,71 +25,92 @@ def main():
     """)
 
     uploaded_file = st.file_uploader("Upload a CIF file", type=["cif"], accept_multiple_files=False)
-    # uploaded_file = "ABABEM.cif"
+    
     if uploaded_file is not None:
         try:
-            with open(uploaded_file.name, "wb") as f:
-                f.write(uploaded_file.getbuffer())
             filename = str(uploaded_file.name)
-            # Read the CIF file using ASE
-            atoms = read(filename, format="cif")
-            cif = ReadCif(filename)
-            cif_data = cif.first_block()
-            if "_diffrn_ambient_temperature" in cif_data.keys():
-                temperature = float(cif_data["_diffrn_ambient_temperature"])
-            else:
-                raise ValueError("Temperature not found in the CIF file. \
-                                    Please provide a temperature in the field _diffrn_ambient_temperature from the CIF file.")
-            st.success("CIF file successfully read.")
-
-            data = Data()
-            data.x = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int32)
+            file = BytesIO(uploaded_file.getbuffer())
+            cif = ReadCif(file)
+            print(cif.keys())
+            if len(cif.keys())>1:
+                st.markdown("Found " + str(len(cif.keys())) + " blocks in the CIF file. We will process all of them and export as separate CIF files.")
+            for key in cif.keys():
+                print(key)
+                # print(cif[key])
+                block = "data_"+str(key)+"\n"+ cif[key].printsection()
+                atoms = read(StringIO(block), format="cif")
+                print("atoms")
+                print(atoms)
+            # atoms = read(atoms_2, format="cif")
+            # with open(uploaded_file.name, "wb") as f:
+            #     f.write(uploaded_file.getbuffer())
+            # filename = str(uploaded_file.name)
+            # # Read the CIF file using ASE
+            # atoms = read(filename, format="cif")
+            # cif = ReadCif(filename)
+            # print(cif.keys())
+            # print(len(atoms))
+            # # st.markdown(cif)
+            # cif_data = cif
+            # st.markdown(f"### CIF file: {filename}")
+            # temperature = 100 
+            # if "_diffrn_ambient_temperature" in cif_data.keys():
+            #     temperature = float(cif_data["_diffrn_ambient_temperature"].split("(")[0])
+            # elif "_cell_measurement_temperature" in cif_data.keys():
+            #     temperature = float(cif_data["_cell_measurement_temperature"].split("(")[0])
+            # else:
+            #     raise ValueError("Temperature not found in the CIF file. \
+            #                         Please provide a temperature in the field _diffrn_ambient_temperature o in the field _cell_measurement_temperature from the CIF file.")
+            # st.success("CIF file successfully read.")
+
+            # data = Data()
+            # data.x = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int32)
             
-            data.pos = torch.tensor(atoms.positions, dtype=torch.float32)
-            data.temperature_og = torch.tensor([temperature], dtype=torch.float32)
-            data.temperature = (data.temperature_og - MEAN_TEMP) / STD_TEMP
-            data.cell = torch.tensor(atoms.cell.array, dtype=torch.float32).unsqueeze(0)
+            # data.pos = torch.tensor(atoms.positions, dtype=torch.float32)
+            # data.temperature_og = torch.tensor([temperature], dtype=torch.float32)
+            # data.temperature = (data.temperature_og - MEAN_TEMP) / STD_TEMP
+            # data.cell = torch.tensor(atoms.cell.array, dtype=torch.float32).unsqueeze(0)
 
-            data.pbc = torch.tensor([True, True, True])
-            data.natoms = len(atoms)
+            # data.pbc = torch.tensor([True, True, True])
+            # data.natoms = len(atoms)
 
-            del atoms
-            gc.collect()
-            batch = Batch.from_data_list([data])
+            # del atoms
+            # gc.collect()
+            # batch = Batch.from_data_list([data])
             
 
-            edge_index, _, _, edge_attr = radius_graph_pbc(batch, 5.0, 64)
-            del batch
-            gc.collect()
-            data.cart_dist = torch.norm(edge_attr, dim=-1)
-            data.cart_dir = torch.nn.functional.normalize(edge_attr, dim=-1)
-            data.edge_index = edge_index
-            data.non_H_mask = data.x != 1
-            delattr(data, "pbc")
-            delattr(data, "natoms")
-            batch = Batch.from_data_list([data])
-            del data, edge_index, edge_attr
-            gc.collect()
-
-            st.success("Graph successfully created.")
-
-            process_data(batch, model)
-            st.success("ADPs successfully predicted.")
+            # edge_index, _, _, edge_attr = radius_graph_pbc(batch, 5.0, 64)
+            # del batch
+            # gc.collect()
+            # data.cart_dist = torch.norm(edge_attr, dim=-1)
+            # data.cart_dir = torch.nn.functional.normalize(edge_attr, dim=-1)
+            # data.edge_index = edge_index
+            # data.non_H_mask = data.x != 1
+            # delattr(data, "pbc")
+            # delattr(data, "natoms")
+            # batch = Batch.from_data_list([data])
+            # del data, edge_index, edge_attr
+            # gc.collect()
+
+            # st.success("Graph successfully created.")
+
+            # process_data(batch, model)
+            # st.success("ADPs successfully predicted.")
             
-            # Create a download button for the processed CIF file
-            with open("output.cif", "r") as f:
-                cif_contents = f.read()
+            # # Create a download button for the processed CIF file
+            # with open("output.cif", "r") as f:
+            #     cif_contents = f.read()
             
-            st.download_button(
-                label="Download processed CIF file",
-                data=cif_contents,
-                file_name="output.cif",
-                mime="text/plain"
-            )
-
-            os.remove("output.cif")
-            os.remove(filename)
-            gc.collect()
+            # st.download_button(
+            #     label="Download processed CIF file",
+            #     data=cif_contents,
+            #     file_name="output.cif",
+            #     mime="text/plain"
+            # )
+
+            # os.remove("output.cif")
+            # os.remove(filename)
+            # gc.collect()
         except Exception as e:
             st.error(f"An error occurred while reading the CIF file: {e}")