Add carbon emissions info and summary (stfc#258)

* Add task emissions to Atoms info
* Add carbon emissions summary

janus_core/calculations/descriptors.py

@@ -120,7 +120,7 @@ def run(self) -> None:
         """Calculate descriptors for structure(s)."""
         if self.logger:
             self.logger.info("Starting descriptors calculation")
-            self.tracker.start()
+            self.tracker.start_task("Descriptors")
 
         if isinstance(self.struct, Sequence):
             for struct in self.struct:
@@ -129,6 +129,12 @@ def run(self) -> None:
             self._calc_descriptors(self.struct)
 
         if self.logger:
+            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")
 

janus_core/calculations/eos.py

@@ -247,7 +247,8 @@ def run(self) -> EoSResults:
         bulk_modulus *= 1.0e24 / kJ
 
         if self.logger:
-            self.tracker.stop_task()
+            emissions = self.tracker.stop_task().emissions
+            self.struct.info["emissions"] = emissions
             self.tracker.stop()
             self.logger.info("Equation of state fitting complete")
 
@@ -302,5 +303,6 @@ def _calc_volumes_energies(self) -> None:
             )
 
         if self.logger:
-            self.tracker.stop_task()
+            emissions = self.tracker.stop_task().emissions
+            self.struct.info["emissions"] = emissions
             self.logger.info("Calculations for configurations complete")

janus_core/calculations/geom_opt.py

@@ -255,7 +255,7 @@ def run(self) -> None:
 
         if self.logger:
             self.logger.info("Starting geometry optimization")
-            self.tracker.start()
+            self.tracker.start_task("Geometry optimization")
 
         converged = self.dyn.run(fmax=self.fmax, steps=self.steps)
 
@@ -296,5 +296,7 @@ def run(self) -> None:
             )
 
         if self.logger:
+            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

@@ -840,7 +840,8 @@ def _run_dynamics(self) -> None:
 
             if self.logger:
                 self.logger.info("Temperature ramp complete at %sK", temps[-1])
-                self.tracker.stop_task()
+                emissions = self.tracker.stop_task().emissions
+                self.struct.info["emissions"] = emissions
 
         # Run MD
         if self.steps > 0:
@@ -856,7 +857,8 @@ def _run_dynamics(self) -> None:
             self._write_final_state()
             self.created_final_file = True
             if self.logger:
-                self.tracker.stop_task()
+                emissions = self.tracker.stop_task().emissions
+                self.struct.info["emissions"] = emissions
                 self.tracker.stop()
                 self.logger.info("Molecular dynamics simulation complete")
 

janus_core/calculations/phonons.py

@@ -317,7 +317,8 @@ def calc_force_constants(
             self.results["phonon"].symmetrize_force_constants(level=1)
 
         if self.logger:
-            self.tracker.stop_task()
+            emissions = self.tracker.stop_task().emissions
+            self.struct.info["emissions"] = emissions
             self.tracker.flush()
             self.logger.info("Phonons calculation complete")
 
@@ -470,7 +471,8 @@ def calc_thermal_props(
         ].get_thermal_properties_dict()
 
         if self.logger:
-            self.tracker.stop_task()
+            emissions = self.tracker.stop_task().emissions
+            self.struct.info["emissions"] = emissions
             self.tracker.flush()
             self.logger.info("Thermal properties calculation complete")
 
@@ -540,7 +542,8 @@ def calc_dos(
         self.results["phonon"].run_total_dos()
 
         if self.logger:
-            self.tracker.stop_task()
+            emissions = self.tracker.stop_task().emissions
+            self.struct.info["emissions"] = emissions
             self.tracker.flush()
             self.logger.info("DOS calculation complete")
 
@@ -643,7 +646,8 @@ def calc_pdos(
         self.results["phonon"].run_projected_dos()
 
         if self.logger:
-            self.tracker.stop_task()
+            emissions = self.tracker.stop_task().emissions
+            self.struct.info["emissions"] = emissions
             self.tracker.flush()
             self.logger.info("PDOS calculation complete")
 

janus_core/calculations/single_point.py

@@ -331,7 +331,12 @@ def run(self) -> CalcResults:
             self.results["stress"] = self._get_stress()
 
         if self.logger:
-            self.tracker.stop_task()
+            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("Single point calculation complete")
 

janus_core/cli/descriptors.py

@@ -20,6 +20,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import (
+    carbon_summary,
     check_config,
     end_summary,
     parse_typer_dicts,
@@ -168,5 +169,7 @@ def descriptors(
     descript = Descriptors(**descriptors_kwargs)
     descript.run()
 
+    carbon_summary(summary=summary, log=log)
+
     # Time after optimization has finished
     end_summary(summary)

janus_core/cli/eos.py

@@ -21,6 +21,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import (
+    carbon_summary,
     check_config,
     end_summary,
     parse_typer_dicts,
@@ -206,5 +207,7 @@ def eos(
     equation_of_state = EoS(**eos_kwargs)
     equation_of_state.run()
 
+    carbon_summary(summary=summary, log=log)
+
     # Time after calculations have finished
     end_summary(summary)

janus_core/cli/geomopt.py

@@ -21,6 +21,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import (
+    carbon_summary,
     check_config,
     end_summary,
     parse_typer_dicts,
@@ -288,5 +289,7 @@ def geomopt(
     optimizer = GeomOpt(**optimize_kwargs)
     optimizer.run()
 
+    carbon_summary(summary=summary, log=log)
+
     # Time after optimization has finished
     end_summary(summary)

janus_core/cli/md.py

@@ -23,6 +23,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import (
+    carbon_summary,
     check_config,
     end_summary,
     parse_typer_dicts,
@@ -425,5 +426,7 @@ def md(
     # Run molecular dynamics
     dyn.run()
 
+    carbon_summary(summary=summary, log=log)
+
     # Save time after simulation has finished
     end_summary(summary=summary)

janus_core/cli/phonons.py

@@ -20,6 +20,7 @@
     Summary,
 )
 from janus_core.cli.utils import (
+    carbon_summary,
     check_config,
     end_summary,
     parse_typer_dicts,
@@ -287,5 +288,7 @@ def phonons(
     phonon = Phonons(**phonons_kwargs)
     phonon.run()
 
+    carbon_summary(summary=summary, log=log)
+
     # Time after calculations have finished
     end_summary(summary)

janus_core/cli/singlepoint.py

@@ -19,6 +19,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import (
+    carbon_summary,
     check_config,
     end_summary,
     parse_typer_dicts,
@@ -154,5 +155,7 @@ def singlepoint(
     # Run singlepoint calculation
     s_point.run()
 
+    carbon_summary(summary=summary, log=log)
+
     # Save time after simulation has finished
     end_summary(summary)

janus_core/cli/utils.py

@@ -113,6 +113,30 @@ def start_summary(*, command: str, summary: Path, inputs: dict) -> None:
         yaml.dump(save_info, outfile, default_flow_style=False)
 
 
+def carbon_summary(*, summary: Path, log: Path) -> None:
+    """
+    Calculate and write carbon tracking summary.
+
+    Parameters
+    ----------
+    summary : Path
+        Path to summary file being saved.
+    log : Path
+        Path to log file with carbon emissions saved.
+    """
+    with open(log, encoding="utf8") as file:
+        logs = yaml.safe_load(file)
+
+    emissions = sum(
+        lg["message"]["emissions"]
+        for lg in logs
+        if isinstance(lg["message"], dict) and "emissions" in lg["message"]
+    )
+
+    with open(summary, "a", encoding="utf8") as outfile:
+        yaml.dump({"emissions": emissions}, outfile, default_flow_style=False)
+
+
 def end_summary(summary: Path) -> None:
     """
     Write final time to summary and close.

tests/test_descriptors.py

@@ -64,3 +64,27 @@ def test_calc_per_element(tmp_path):
     assert atoms.info["mace_descriptor"] == pytest.approx(-0.005626419559511429)
     assert atoms.info["mace_Cl_descriptor"] == pytest.approx(-0.009215340539869301)
     assert atoms.info["mace_Na_descriptor"] == pytest.approx(-0.0020374985791535563)
+
+
+def test_logging(tmp_path):
+    """Test attaching logger to Descriptors and emissions are saved to info."""
+    log_file = tmp_path / "descriptors.log"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace_mp",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    descriptors = Descriptors(
+        single_point.struct,
+        calc_per_element=True,
+        log_kwargs={"filename": log_file},
+    )
+
+    assert "emissions" not in single_point.struct.info
+
+    descriptors.run()
+
+    assert log_file.exists()
+    assert single_point.struct.info["emissions"] > 0

tests/test_descriptors_cli.py

@@ -5,6 +5,7 @@
 from ase.io import read
 import pytest
 from typer.testing import CliRunner
+import yaml
 
 from janus_core.cli.janus import app
 from tests.utils import assert_log_contains, strip_ansi_codes
@@ -59,6 +60,20 @@ def test_descriptors(tmp_path):
         ],
     )
 
+    # Read descriptors summary file
+    assert summary_path.exists()
+    with open(summary_path, encoding="utf8") as file:
+        descriptors_summary = yaml.safe_load(file)
+
+    assert "command" in descriptors_summary
+    assert "janus descriptors" in descriptors_summary["command"]
+    assert "start_time" in descriptors_summary
+    assert "inputs" in descriptors_summary
+    assert "end_time" in descriptors_summary
+
+    assert "emissions" in descriptors_summary
+    assert descriptors_summary["emissions"] > 0
+
 
 def test_calc_per_element(tmp_path):
     """Test calculating MLIP descriptors for each element."""

tests/test_eos.py

@@ -108,3 +108,27 @@ def test_invalid_struct():
         EoS(
             "structure",
         )
+
+
+def test_logging(tmp_path):
+    """Test attaching logger to EoS and emissions are saved to info."""
+    log_file = tmp_path / "eos.log"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace_mp",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    eos = EoS(
+        single_point.struct,
+        file_prefix=tmp_path / "NaCl",
+        log_kwargs={"filename": log_file},
+    )
+
+    assert "emissions" not in single_point.struct.info
+
+    eos.run()
+
+    assert log_file.exists()
+    assert single_point.struct.info["emissions"] > 0

tests/test_eos_cli.py

@@ -5,6 +5,7 @@
 from ase.io import read
 import pytest
 from typer.testing import CliRunner
+import yaml
 
 from janus_core.cli.janus import app
 from tests.utils import assert_log_contains, strip_ansi_codes
@@ -72,6 +73,20 @@ def test_eos(tmp_path):
         excludes=["Minimising lattice scalar = 1.0"],
     )
 
+    # Read eos summary file
+    assert summary_path.exists()
+    with open(summary_path, encoding="utf8") as file:
+        eos_summary = yaml.safe_load(file)
+
+    assert "command" in eos_summary
+    assert "janus eos" in eos_summary["command"]
+    assert "start_time" in eos_summary
+    assert "inputs" in eos_summary
+    assert "end_time" in eos_summary
+
+    assert "emissions" in eos_summary
+    assert eos_summary["emissions"] > 0
+
 
 def test_setting_lattice(tmp_path):
     """Test setting the lattice constants."""

tests/test_geom_opt.py

@@ -333,3 +333,25 @@ def test_invalid_struct():
             fmax=0.001,
             optimizer="test",
         )
+
+
+def test_logging(tmp_path):
+    """Test attaching logger to GeomOpt and emissions are saved to info."""
+    log_file = tmp_path / "geomopt.log"
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace_mp",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    assert "emissions" not in single_point.struct.info
+
+    optimizer = GeomOpt(
+        single_point.struct,
+        log_kwargs={"filename": log_file},
+    )
+    optimizer.run()
+
+    assert log_file.exists()
+    assert "emissions" in single_point.struct.info
+    assert single_point.struct.info["emissions"] > 0