Merge pull request #195 from ICAMS/fix_comp

Fix comp

Author: srmnitc

calphy/input.py

@@ -181,10 +181,17 @@ class MeltingTemperature(BaseModel, title="Input options for melting temperature
     step: Annotated[int, Field(default=200, ge=20)]
     attempts: Annotated[int, Field(default=5, ge=1)]
 
-class MaterialsProject(BaseModel, title='Input options for materials project'):
+
+class MaterialsProject(BaseModel, title="Input options for materials project"):
     api_key: Annotated[str, Field(default="", exclude=True)]
     conventional: Annotated[bool, Field(default=True)]
-    target_natoms: Annotated[int, Field(default=1500, description='The structure parsed from materials project would be repeated to approximately this value')]
+    target_natoms: Annotated[
+        int,
+        Field(
+            default=1500,
+            description="The structure parsed from materials project would be repeated to approximately this value",
+        ),
+    ]
 
     @field_validator("api_key", mode="after")
     def resolve_api_key(cls, v: str) -> str:
@@ -198,6 +205,7 @@ def resolve_api_key(cls, v: str) -> str:
             )
         return value
 
+
 class Calculation(BaseModel, title="Main input class"):
     monte_carlo: Optional[MonteCarlo] = MonteCarlo()
     composition_scaling: Optional[CompositionScaling] = CompositionScaling()
@@ -515,40 +523,58 @@ def _validate_all(self) -> "Input":
             self._original_lattice = self.lattice.lower()
             write_structure_file = True
 
-        elif self.lattice.split('-')[0] == 'mp':
-            #confirm here that API key exists
+        elif self.lattice.split("-")[0] == "mp":
+            # confirm here that API key exists
             if not self.materials_project.api_key:
-                raise ValueError('could not find API KEY, pls set it.')
-            #now we need to fetch the structure
+                raise ValueError("could not find API KEY, pls set it.")
+            # now we need to fetch the structure
             try:
                 from mp_api.client import MPRester
             except ImportError:
-                raise ImportError('Could not import mp_api, make sure you install mp_api package!')
-            #now all good
+                raise ImportError(
+                    "Could not import mp_api, make sure you install mp_api package!"
+                )
+            # now all good
             rest = {
-                    "use_document_model": False,
-                    "include_user_agent": True,
-                    "api_key": self.materials_project.api_key, 
-                }
+                "use_document_model": False,
+                "include_user_agent": True,
+                "api_key": self.materials_project.api_key,
+            }
             with MPRester(**rest) as mpr:
                 docs = mpr.materials.summary.search(material_ids=[self.lattice])
 
             structures = []
             for doc in docs:
-                struct = doc['structure']
+                struct = doc["structure"]
                 if self.materials_project.conventional:
                     aseatoms = struct.to_conventional().to_ase_atoms()
                 else:
                     aseatoms = struct.to_primitive().to_ase_atoms()
                 structures.append(aseatoms)
             structure = structures[0]
-            
+
             if np.prod(self.repeat) == 1:
-                x = int(np.ceil((self.materials_project.target_natoms/len(structure))**(1/3)))
+                x = int(
+                    np.ceil(
+                        (self.materials_project.target_natoms / len(structure))
+                        ** (1 / 3)
+                    )
+                )
                 structure = structure.repeat(x)
             else:
                 structure = structure.repeat(self.repeat)
 
+            # extract composition
+            types, typecounts = np.unique(
+                structure.get_chemical_symbols(), return_counts=True
+            )
+
+            for c, t in enumerate(types):
+                self._element_dict[t]["count"] = typecounts[c]
+                self._element_dict[t]["composition"] = typecounts[c] / np.sum(
+                    typecounts
+                )
+
             self._natoms = len(structure)
             self._original_lattice = self.lattice.lower()
             write_structure_file = True

calphy/postprocessing.py

@@ -30,15 +30,15 @@ def read_report(folder):
 
 def _extract_error(errfile):
     error_code = None
-    if os.path.exists(errfile):
-        with open(errfile, 'r') as fin:
-            for line in fin:
-                if 'calphy.errors' in line:
-                    break
-            try:
-                error_code = line.split(':')[0].split('.')[-1]
-            except:
-                pass
+    try:
+        if os.path.exists(errfile):
+            with open(errfile, 'r') as fin:
+                for line in fin:
+                    if 'calphy.errors' in line:
+                        break
+                    error_code = line.split(':')[0].split('.')[-1]
+    except:
+        pass
     return error_code
 
 def gather_results(mainfolder, reduce_composition=True,