Make carbon tracking optional (#343)

Author: ElliottKasoar

janus_core/calculations/base.py

@@ -47,6 +47,8 @@ class BaseCalculation(FileNameMixin):
         Whether to attach a logger. Default is False.
     log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_logger`. Default is {}.
+    track_carbon : bool
+        Whether to track carbon emissions of calculation. Default is True.
     tracker_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_tracker`. Default is {}.
     file_prefix : Optional[PathLike]
@@ -79,6 +81,7 @@ def __init__(
         set_calc: Optional[bool] = None,
         attach_logger: bool = False,
         log_kwargs: Optional[dict[str, Any]] = None,
+        track_carbon: bool = True,
         tracker_kwargs: Optional[dict[str, Any]] = None,
         file_prefix: Optional[PathLike] = None,
         additional_prefix: Optional[str] = None,
@@ -115,6 +118,8 @@ def __init__(
             Whether to attach a logger. Default is False.
         log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_logger`. Default is {}.
+        track_carbon : bool
+            Whether to track carbon emissions of calculation. Default is True.
         tracker_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_tracker`. Default is {}.
         file_prefix : Optional[PathLike]
@@ -137,6 +142,7 @@ def __init__(
         self.read_kwargs = read_kwargs
         self.calc_kwargs = calc_kwargs
         self.log_kwargs = log_kwargs
+        self.track_carbon = track_carbon
         self.tracker_kwargs = tracker_kwargs
 
         if not self.model_path and "model_path" in self.calc_kwargs:
@@ -179,4 +185,6 @@ def __init__(
 
         self.log_kwargs.setdefault("name", calc_name)
         self.logger = config_logger(**self.log_kwargs)
-        self.tracker = config_tracker(self.logger, **self.tracker_kwargs)
+        self.tracker = config_tracker(
+            self.logger, self.track_carbon, **self.tracker_kwargs
+        )

janus_core/calculations/descriptors.py

@@ -47,6 +47,8 @@ class Descriptors(BaseCalculation):
         Whether to attach a logger. Default is False.
     log_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_logger`. Default is {}.
+    track_carbon : bool
+        Whether to track carbon emissions of calculation. Default is True.
     tracker_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_tracker`. Default is {}.
     invariants_only : bool
@@ -79,6 +81,7 @@ def __init__(
         set_calc: Optional[bool] = None,
         attach_logger: bool = False,
         log_kwargs: Optional[dict[str, Any]] = None,
+        track_carbon: bool = True,
         tracker_kwargs: Optional[dict[str, Any]] = None,
         invariants_only: bool = True,
         calc_per_element: bool = False,
@@ -114,6 +117,8 @@ def __init__(
             Whether to attach a logger. Default is False.
         log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_logger`. Default is {}.
+        track_carbon : bool
+            Whether to track carbon emissions of calculation. Default is True.
         tracker_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_tracker`. Default is {}.
         invariants_only : bool
@@ -153,6 +158,7 @@ def __init__(
             set_calc=set_calc,
             attach_logger=attach_logger,
             log_kwargs=log_kwargs,
+            track_carbon=track_carbon,
             tracker_kwargs=tracker_kwargs,
         )
 
@@ -182,6 +188,7 @@ def run(self) -> None:
             self.logger.info("invariants_only: %s", self.invariants_only)
             self.logger.info("calc_per_element: %s", self.calc_per_element)
             self.logger.info("calc_per_atom: %s", self.calc_per_atom)
+        if self.tracker:
             self.tracker.start_task("Descriptors")
 
         if isinstance(self.struct, Sequence):
@@ -191,14 +198,15 @@ def run(self) -> None:
             self._calc_descriptors(self.struct)
 
         if self.logger:
+            self.logger.info("Descriptors calculation complete")
+        if self.tracker:
             emissions = self.tracker.stop_task().emissions
             if isinstance(self.struct, Sequence):
                 for image in self.struct:
                     image.info["emissions"] = emissions
             else:
                 self.struct.info["emissions"] = emissions
             self.tracker.stop()
-            self.logger.info("Descriptors calculation complete")
 
         output_structs(
             self.struct,

janus_core/calculations/eos.py

@@ -51,6 +51,8 @@ class EoS(BaseCalculation):
         Whether to attach a logger. Default is False.
     log_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_logger`. Default is {}.
+    track_carbon : bool
+        Whether to track carbon emissions of calculation. Default is True.
     tracker_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_tracker`. Default is {}.
     min_volume : float
@@ -112,6 +114,7 @@ def __init__(
         set_calc: Optional[bool] = None,
         attach_logger: bool = False,
         log_kwargs: Optional[dict[str, Any]] = None,
+        track_carbon: bool = True,
         tracker_kwargs: Optional[dict[str, Any]] = None,
         min_volume: float = 0.95,
         max_volume: float = 1.05,
@@ -155,6 +158,8 @@ def __init__(
             Whether to attach a logger. Default is False.
         log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_logger`. Default is {}.
+        track_carbon : bool
+            Whether to track carbon emissions of calculation. Default is True.
         tracker_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_tracker`. Default is {}.
         min_volume : float
@@ -241,6 +246,7 @@ def __init__(
             set_calc=set_calc,
             attach_logger=attach_logger,
             log_kwargs=log_kwargs,
+            track_carbon=track_carbon,
             tracker_kwargs=tracker_kwargs,
             file_prefix=file_prefix,
         )
@@ -313,17 +319,19 @@ def run(self) -> EoSResults:
 
         if self.logger:
             self.logger.info("Starting of fitting equation of state")
+        if self.tracker:
             self.tracker.start_task("Fit EoS")
 
         v_0, e_0, bulk_modulus = eos.fit()
         # transform bulk modulus unit in GPa
         bulk_modulus *= 1.0e24 / kJ
 
         if self.logger:
+            self.logger.info("Equation of state fitting complete")
+        if self.tracker:
             emissions = self.tracker.stop_task().emissions
             self.struct.info["emissions"] = emissions
             self.tracker.stop()
-            self.logger.info("Equation of state fitting complete")
 
         if self.write_results:
             with open(f"{self.file_prefix}-eos-fit.dat", "w", encoding="utf8") as out:
@@ -346,6 +354,7 @@ def _calc_volumes_energies(self) -> None:
         """Calculate volumes and energies for all lattice constants."""
         if self.logger:
             self.logger.info("Starting calculations for configurations")
+        if self.tracker:
             self.tracker.start_task("Calculate configurations")
 
         cell = self.struct.get_cell()
@@ -380,6 +389,7 @@ def _calc_volumes_energies(self) -> None:
             )
 
         if self.logger:
+            self.logger.info("Calculations for configurations complete")
+        if self.tracker:
             emissions = self.tracker.stop_task().emissions
             self.struct.info["emissions"] = emissions
-            self.logger.info("Calculations for configurations complete")

janus_core/calculations/geom_opt.py

@@ -51,6 +51,8 @@ class GeomOpt(BaseCalculation):
         Whether to attach a logger. Default is False.
     log_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_logger`. Default is {}.
+    track_carbon : bool
+        Whether to track carbon emissions of calculation. Default is True.
     tracker_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_tracker`. Default is {}.
     fmax : float
@@ -104,6 +106,7 @@ def __init__(
         set_calc: Optional[bool] = None,
         attach_logger: bool = False,
         log_kwargs: Optional[dict[str, Any]] = None,
+        track_carbon: bool = True,
         tracker_kwargs: Optional[dict[str, Any]] = None,
         fmax: float = 0.1,
         steps: int = 1000,
@@ -146,6 +149,8 @@ def __init__(
             Whether to attach a logger. Default is False.
         log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_logger`. Default is {}.
+        track_carbon : bool
+            Whether to track carbon emissions of calculation. Default is True.
         tracker_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_tracker`. Default is {}.
         fmax : float
@@ -224,6 +229,7 @@ def __init__(
             set_calc=set_calc,
             attach_logger=attach_logger,
             log_kwargs=log_kwargs,
+            track_carbon=track_carbon,
             tracker_kwargs=tracker_kwargs,
         )
 
@@ -292,6 +298,7 @@ def run(self) -> None:
 
         if self.logger:
             self.logger.info("Starting geometry optimization")
+        if self.tracker:
             self.tracker.start_task("Geometry optimization")
 
         converged = self.dyn.run(fmax=self.fmax, steps=self.steps)
@@ -351,7 +358,8 @@ def run(self) -> None:
             )
 
         if self.logger:
+            self.logger.info("Geometry optimization complete")
+        if self.tracker:
             emissions = self.tracker.stop_task().emissions
             self.struct.info["emissions"] = emissions
             self.tracker.stop()
-            self.logger.info("Geometry optimization complete")

janus_core/calculations/md.py

@@ -76,6 +76,8 @@ class MolecularDynamics(BaseCalculation):
         Whether to attach a logger. Default is False.
     log_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_logger`. Default is {}.
+    track_carbon : bool
+        Whether to track carbon emissions of calculation. Default is True.
     tracker_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to `config_tracker`. Default is {}.
     struct : Atoms
@@ -187,6 +189,7 @@ def __init__(
         set_calc: Optional[bool] = None,
         attach_logger: bool = False,
         log_kwargs: Optional[dict[str, Any]] = None,
+        track_carbon: bool = True,
         tracker_kwargs: Optional[dict[str, Any]] = None,
         ensemble: Optional[Ensembles] = None,
         steps: int = 0,
@@ -250,6 +253,8 @@ def __init__(
             Whether to attach a logger. Default is False.
         log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_logger`. Default is {}.
+        track_carbon : bool
+            Whether to track carbon emissions of calculation. Default is True.
         tracker_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to `config_tracker`. Default is {}.
         ensemble : Ensembles
@@ -448,6 +453,7 @@ def __init__(
             set_calc=set_calc,
             attach_logger=attach_logger,
             log_kwargs=log_kwargs,
+            track_carbon=track_carbon,
             tracker_kwargs=tracker_kwargs,
             file_prefix=file_prefix,
             additional_prefix=self.ensemble,
@@ -1070,6 +1076,7 @@ def _run_dynamics(self) -> None:
 
             if self.logger:
                 self.logger.info("Beginning temperature ramp at %sK", temps[0])
+            if self.tracker:
                 self.tracker.start_task("Temperature ramp")
 
             for temp in temps:
@@ -1087,13 +1094,15 @@ def _run_dynamics(self) -> None:
 
             if self.logger:
                 self.logger.info("Temperature ramp complete at %sK", temps[-1])
+            if self.tracker:
                 emissions = self.tracker.stop_task().emissions
                 self.struct.info["emissions"] = emissions
 
         # Run MD
         if self.steps > 0:
             if self.logger:
                 self.logger.info("Starting molecular dynamics simulation")
+            if self.tracker:
                 self.tracker.start_task("Molecular dynamics")
             self.temp = md_temp
             if self.ramp_temp:
@@ -1104,10 +1113,11 @@ def _run_dynamics(self) -> None:
             self._write_final_state()
             self.created_final_file = True
             if self.logger:
+                self.logger.info("Molecular dynamics simulation complete")
+            if self.tracker:
                 emissions = self.tracker.stop_task().emissions
                 self.struct.info["emissions"] = emissions
                 self.tracker.stop()
-                self.logger.info("Molecular dynamics simulation complete")
 
 
 class NPT(MolecularDynamics):