From aaf0536f80ac398c09d17199f6187241165a3c88 Mon Sep 17 00:00:00 2001 From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com> Date: Tue, 29 Oct 2024 19:23:48 +0000 Subject: [PATCH] Make carbon tracking optional --- janus_core/calculations/base.py | 10 ++++- janus_core/calculations/descriptors.py | 10 ++++- janus_core/calculations/eos.py | 14 ++++++- janus_core/calculations/geom_opt.py | 10 ++++- janus_core/calculations/md.py | 12 +++++- janus_core/calculations/phonons.py | 22 ++++++++-- janus_core/calculations/single_point.py | 10 ++++- janus_core/cli/descriptors.py | 10 ++++- janus_core/cli/eos.py | 10 ++++- janus_core/cli/geomopt.py | 10 ++++- janus_core/cli/md.py | 10 ++++- janus_core/cli/phonons.py | 10 ++++- janus_core/cli/singlepoint.py | 10 ++++- janus_core/cli/train.py | 16 +++++++- janus_core/helpers/log.py | 5 ++- janus_core/helpers/train.py | 7 +++- tests/test_descriptors_cli.py | 29 ++++++++++++++ tests/test_eos_cli.py | 26 ++++++++++++ tests/test_geomopt_cli.py | 29 ++++++++++++++ tests/test_md_cli.py | 29 ++++++++++++++ tests/test_phonons_cli.py | 25 ++++++++++++ tests/test_singlepoint_cli.py | 31 +++++++++++++++ tests/test_train_cli.py | 53 +++++++++++++++++++++++++ 23 files changed, 377 insertions(+), 21 deletions(-) diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py index e66be99d..cdf4a314 100644 --- a/janus_core/calculations/base.py +++ b/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 + ) diff --git a/janus_core/calculations/descriptors.py b/janus_core/calculations/descriptors.py index 13e6e9a4..9e0bbe4f 100644 --- a/janus_core/calculations/descriptors.py +++ b/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,6 +198,8 @@ 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: @@ -198,7 +207,6 @@ def run(self) -> None: else: self.struct.info["emissions"] = emissions self.tracker.stop() - self.logger.info("Descriptors calculation complete") output_structs( self.struct, diff --git a/janus_core/calculations/eos.py b/janus_core/calculations/eos.py index 7a97b5da..5a332570 100644 --- a/janus_core/calculations/eos.py +++ b/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,6 +319,7 @@ 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() @@ -320,10 +327,11 @@ def run(self) -> EoSResults: 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") diff --git a/janus_core/calculations/geom_opt.py b/janus_core/calculations/geom_opt.py index 350799a5..bb859dd5 100644 --- a/janus_core/calculations/geom_opt.py +++ b/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") diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py index 026d7adf..371ed694 100644 --- a/janus_core/calculations/md.py +++ b/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,6 +1094,7 @@ 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 @@ -1094,6 +1102,7 @@ def _run_dynamics(self) -> None: 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): diff --git a/janus_core/calculations/phonons.py b/janus_core/calculations/phonons.py index c8325fcd..b3262cab 100644 --- a/janus_core/calculations/phonons.py +++ b/janus_core/calculations/phonons.py @@ -52,6 +52,8 @@ class Phonons(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 {}. calcs : Optional[MaybeSequence[PhononCalcs]] @@ -138,6 +140,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, calcs: MaybeSequence[PhononCalcs] = (), supercell: MaybeList[int] = 2, @@ -184,6 +187,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 {}. calcs : Optional[MaybeSequence[PhononCalcs]] @@ -264,6 +269,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, ) @@ -359,6 +365,7 @@ def calc_force_constants( if self.logger: self.logger.info("Starting phonons calculation") + if self.tracker: self.tracker.start_task("Phonon calculation") cell = self._ASE_to_PhonopyAtoms(self.struct) @@ -390,10 +397,11 @@ def calc_force_constants( self.results["phonon"].symmetrize_force_constants(level=1) if self.logger: + self.logger.info("Phonons calculation complete") + if self.tracker: emissions = self.tracker.stop_task().emissions self.struct.info["emissions"] = emissions self.tracker.flush() - self.logger.info("Phonons calculation complete") if write_force_consts: self.write_force_constants(**kwargs) @@ -541,6 +549,7 @@ def calc_thermal_props( if self.logger: self.logger.info("Starting thermal properties calculation") + if self.tracker: self.tracker.start_task("Thermal calculation") self.results["phonon"].run_mesh(mesh) @@ -552,10 +561,11 @@ def calc_thermal_props( ].get_thermal_properties_dict() if self.logger: + self.logger.info("Thermal properties calculation complete") + if self.tracker: emissions = self.tracker.stop_task().emissions self.struct.info["emissions"] = emissions self.tracker.flush() - self.logger.info("Thermal properties calculation complete") if write_thermal: self.write_thermal_props(**kwargs) @@ -620,16 +630,18 @@ def calc_dos( if self.logger: self.logger.info("Starting DOS calculation") + if self.tracker: self.tracker.start_task("DOS calculation") self.results["phonon"].run_mesh(mesh) self.results["phonon"].run_total_dos() if self.logger: + self.logger.info("DOS calculation complete") + if self.tracker: emissions = self.tracker.stop_task().emissions self.struct.info["emissions"] = emissions self.tracker.flush() - self.logger.info("DOS calculation complete") if write_dos: self.write_dos(**kwargs) @@ -725,6 +737,7 @@ def calc_pdos( if self.logger: self.logger.info("Starting PDOS calculation") + if self.tracker: self.tracker.start_task("PDOS calculation") self.results["phonon"].run_mesh( @@ -733,10 +746,11 @@ def calc_pdos( self.results["phonon"].run_projected_dos() if self.logger: + self.logger.info("PDOS calculation complete") + if self.tracker: emissions = self.tracker.stop_task().emissions self.struct.info["emissions"] = emissions self.tracker.flush() - self.logger.info("PDOS calculation complete") if write_pdos: self.write_pdos(**kwargs) diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py index d7ae718e..5aa7c29a 100644 --- a/janus_core/calculations/single_point.py +++ b/janus_core/calculations/single_point.py @@ -51,6 +51,8 @@ class SinglePoint(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 {}. properties : MaybeSequence[Properties] @@ -86,6 +88,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, properties: MaybeSequence[Properties] = (), write_results: bool = False, @@ -120,6 +123,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 {}. properties : MaybeSequence[Properties] @@ -154,6 +159,7 @@ def __init__( set_calc=set_calc, attach_logger=attach_logger, log_kwargs=log_kwargs, + track_carbon=track_carbon, tracker_kwargs=tracker_kwargs, ) @@ -308,6 +314,7 @@ def run(self) -> CalcResults: if self.logger: self.logger.info("Starting single point calculation") + if self.tracker: self.tracker.start_task("Single point") if "energy" in self.properties: @@ -320,6 +327,8 @@ def run(self) -> CalcResults: self.results["hessian"] = self._get_hessian() if self.logger: + self.logger.info("Single point calculation complete") + if self.tracker: emissions = self.tracker.stop_task().emissions if isinstance(self.struct, Sequence): for image in self.struct: @@ -327,7 +336,6 @@ def run(self) -> CalcResults: else: self.struct.info["emissions"] = emissions self.tracker.stop() - self.logger.info("Single point calculation complete") output_structs( self.struct, diff --git a/janus_core/cli/descriptors.py b/janus_core/cli/descriptors.py index 30343d03..fa54d1a0 100644 --- a/janus_core/cli/descriptors.py +++ b/janus_core/cli/descriptors.py @@ -56,6 +56,9 @@ def descriptors( calc_kwargs: CalcKwargs = None, write_kwargs: WriteKwargs = None, log: LogPath = None, + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Summary = None, ): """ @@ -92,6 +95,9 @@ def descriptors( Keyword arguments to pass to ase.io.write when saving results. Default is {}. log : Optional[Path] Path to write logs to. Default is inferred from the name of the structure file. + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is inferred from the name of the structure file. @@ -137,6 +143,7 @@ def descriptors( "calc_kwargs": calc_kwargs, "attach_logger": True, "log_kwargs": log_kwargs, + "track_carbon": tracker, "invariants_only": invariants_only, "calc_per_element": calc_per_element, "calc_per_atom": calc_per_atom, @@ -176,7 +183,8 @@ def descriptors( descript.run() # Save carbon summary - carbon_summary(summary=summary, log=log) + if tracker: + carbon_summary(summary=summary, log=log) # Time after calculation has finished end_summary(summary) diff --git a/janus_core/cli/eos.py b/janus_core/cli/eos.py index 6f2a8aa2..0aac05e3 100644 --- a/janus_core/cli/eos.py +++ b/janus_core/cli/eos.py @@ -72,6 +72,9 @@ def eos( ), ] = None, log: LogPath = None, + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Summary = None, ): """ @@ -124,6 +127,9 @@ def eos( chemical formula. log : Optional[Path] Path to write logs to. Default is inferred from the name of the structure file. + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is inferred from the name of the structure file. @@ -174,6 +180,7 @@ def eos( "calc_kwargs": calc_kwargs, "attach_logger": True, "log_kwargs": log_kwargs, + "track_carbon": tracker, "min_volume": min_volume, "max_volume": max_volume, "n_volumes": n_volumes, @@ -219,7 +226,8 @@ def eos( equation_of_state.run() # Save carbon summary - carbon_summary(summary=summary, log=log) + if tracker: + carbon_summary(summary=summary, log=log) # Time after calculations have finished end_summary(summary) diff --git a/janus_core/cli/geomopt.py b/janus_core/cli/geomopt.py index c072c80c..bcc7a19d 100644 --- a/janus_core/cli/geomopt.py +++ b/janus_core/cli/geomopt.py @@ -146,6 +146,9 @@ def geomopt( minimize_kwargs: MinimizeKwargs = None, write_kwargs: WriteKwargs = None, log: LogPath = None, + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Summary = None, ): """ @@ -206,6 +209,9 @@ def geomopt( Default is {}. log : Optional[Path] Path to write logs to. Default is inferred from the name of the structure file. + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is inferred from the name of the structure file. @@ -267,6 +273,7 @@ def geomopt( "calc_kwargs": calc_kwargs, "attach_logger": True, "log_kwargs": log_kwargs, + "track_carbon": tracker, "optimizer": optimizer, "fmax": fmax, "steps": steps, @@ -310,7 +317,8 @@ def geomopt( optimizer.run() # Save carbon summary - carbon_summary(summary=summary, log=log) + if tracker: + carbon_summary(summary=summary, log=log) # Time after optimization has finished end_summary(summary) diff --git a/janus_core/cli/md.py b/janus_core/cli/md.py index df6897c1..fc1bddda 100644 --- a/janus_core/cli/md.py +++ b/janus_core/cli/md.py @@ -170,6 +170,9 @@ def md( Option(help="Random seed for numpy.random and random functions."), ] = None, log: LogPath = None, + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Summary = None, ): """ @@ -281,6 +284,9 @@ def md( Default is None. log : Optional[Path] Path to write logs to. Default is inferred from the name of the structure file. + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is inferred from the name of the structure file. @@ -339,6 +345,7 @@ def md( "calc_kwargs": calc_kwargs, "attach_logger": True, "log_kwargs": log_kwargs, + "track_carbon": tracker, "ensemble_kwargs": ensemble_kwargs, "timestep": timestep, "steps": steps, @@ -436,7 +443,8 @@ def md( dyn.run() # Save carbon summary - carbon_summary(summary=summary, log=log) + if tracker: + carbon_summary(summary=summary, log=log) # Save time after simulation has finished end_summary(summary=summary) diff --git a/janus_core/cli/phonons.py b/janus_core/cli/phonons.py index a6163de6..9617477b 100644 --- a/janus_core/cli/phonons.py +++ b/janus_core/cli/phonons.py @@ -100,6 +100,9 @@ def phonons( ), ] = None, log: LogPath = None, + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Summary = None, ): """ @@ -167,6 +170,9 @@ def phonons( chemical formula. log : Optional[Path] Path to write logs to. Default is inferred from the name of the structure file. + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is inferred from the name of the structure file. @@ -224,6 +230,7 @@ def phonons( "calc_kwargs": calc_kwargs, "attach_logger": True, "log_kwargs": log_kwargs, + "track_carbon": tracker, "calcs": calcs, "supercell": supercell, "displacement": displacement, @@ -275,7 +282,8 @@ def phonons( phonon.run() # Save carbon summary - carbon_summary(summary=summary, log=log) + if tracker: + carbon_summary(summary=summary, log=log) # Time after calculations have finished end_summary(summary) diff --git a/janus_core/cli/singlepoint.py b/janus_core/cli/singlepoint.py index dc298e53..97a9361c 100644 --- a/janus_core/cli/singlepoint.py +++ b/janus_core/cli/singlepoint.py @@ -53,6 +53,9 @@ def singlepoint( calc_kwargs: CalcKwargs = None, write_kwargs: WriteKwargs = None, log: LogPath = None, + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Summary = None, ): """ @@ -85,6 +88,9 @@ def singlepoint( Keyword arguments to pass to ase.io.write when saving results. Default is {}. log : Optional[Path] Path to write logs to. Default is inferred from the name of the structure file. + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is inferred from the name of the structure file. @@ -130,6 +136,7 @@ def singlepoint( "calc_kwargs": calc_kwargs, "attach_logger": True, "log_kwargs": log_kwargs, + "track_carbon": tracker, } # Initialise singlepoint structure and calculator @@ -164,7 +171,8 @@ def singlepoint( s_point.run() # Save carbon summary - carbon_summary(summary=summary, log=log) + if tracker: + carbon_summary(summary=summary, log=log) # Save time after simulation has finished end_summary(summary) diff --git a/janus_core/cli/train.py b/janus_core/cli/train.py index 63d694c2..bb15b49d 100644 --- a/janus_core/cli/train.py +++ b/janus_core/cli/train.py @@ -18,6 +18,9 @@ def train( bool, Option(help="Whether to fine-tune a foundational model.") ] = False, log: Annotated[Path, Option(help="Path to save logs to.")] = Path("train-log.yml"), + tracker: Annotated[ + bool, Option(help="Whether to save carbon emissions of calculation") + ] = True, summary: Annotated[ Path, Option( @@ -38,6 +41,9 @@ def train( Whether to fine-tune a foundational model. Default is False. log : Optional[Path] Path to write logs to. Default is Path("train-log.yml"). + tracker : bool + Whether to save carbon emissions of calculation in log file and summary. + Default is True. summary : Optional[Path] Path to save summary of inputs, start/end time, and carbon emissions. Default is Path("train-summary.yml"). @@ -72,12 +78,18 @@ def train( # Save summary information before training begins start_summary(command="train", summary=summary, inputs=inputs) + log_kwargs = {"filemode": "w"} + if log: + log_kwargs["filename"] = log + # Run training run_train( - mlip_config, attach_logger=True, log_kwargs={"filename": log, "filemode": "w"} + mlip_config, attach_logger=True, log_kwargs=log_kwargs, track_carbon=tracker ) - carbon_summary(summary=summary, log=log) + # Save carbon summary + if tracker: + carbon_summary(summary=summary, log=log) # Save time after training has finished end_summary(summary) diff --git a/janus_core/helpers/log.py b/janus_core/helpers/log.py index 8f6aa390..c61ee35e 100644 --- a/janus_core/helpers/log.py +++ b/janus_core/helpers/log.py @@ -153,6 +153,7 @@ def config_logger( def config_tracker( janus_logger: Optional[logging.Logger], + track_carbon: bool = True, *, country_iso_code: str = "GBR", save_to_file: bool = False, @@ -165,6 +166,8 @@ def config_tracker( ---------- janus_logger : Optional[logging.Logger] Logger for codecarbon output. + track_carbon : bool + Whether to track carbon emissions for calculation. Default is True. country_iso_code : str 3 letter ISO Code of the country where the code is running. Default is "GBR". save_to_file : bool @@ -177,7 +180,7 @@ def config_tracker( Optional[OfflineEmissionsTracker] Configured offline codecarbon tracker, if logger is specified. """ - if janus_logger: + if janus_logger and track_carbon: carbon_logger = LoggerOutput(janus_logger) tracker = OfflineEmissionsTracker( country_iso_code=country_iso_code, diff --git a/janus_core/helpers/train.py b/janus_core/helpers/train.py index a5bae45e..6f9f1d20 100644 --- a/janus_core/helpers/train.py +++ b/janus_core/helpers/train.py @@ -43,6 +43,7 @@ def train( req_file_keys: Optional[list[PathLike]] = None, attach_logger: bool = False, log_kwargs: Optional[dict[str, Any]] = None, + track_carbon: bool = True, tracker_kwargs: Optional[dict[str, Any]] = None, ) -> None: """ @@ -62,6 +63,8 @@ def train( 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 {}. """ @@ -80,7 +83,7 @@ def train( log_kwargs.setdefault("filename", "train-log.yml") log_kwargs.setdefault("name", __name__) logger = config_logger(**log_kwargs) - tracker = config_tracker(logger, **tracker_kwargs) + tracker = config_tracker(logger, track_carbon, **tracker_kwargs) if logger and "foundation_model" in options: logger.info("Fine tuning model: %s", options["foundation_model"]) @@ -89,9 +92,11 @@ def train( mlip_args = mace_parser().parse_args(["--config", str(mlip_config)]) if logger: logger.info("Starting training") + if tracker: tracker.start_task("Training") run_train(mlip_args) if logger: logger.info("Training complete") + if tracker: tracker.stop_task() tracker.stop() diff --git a/tests/test_descriptors_cli.py b/tests/test_descriptors_cli.py index fcc62f7d..3215770e 100644 --- a/tests/test_descriptors_cli.py +++ b/tests/test_descriptors_cli.py @@ -239,3 +239,32 @@ def test_per_atom(tmp_path): "calc_per_atom: True", ], ) + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + out_path = tmp_path / "test" / "NaCl-descriptors.extxyz" + log_path = tmp_path / "test.log" + summary_path = tmp_path / "summary.yml" + + result = runner.invoke( + app, + [ + "descriptors", + "--struct", + DATA_PATH / "NaCl.cif", + "--out", + out_path, + "--log", + log_path, + "--no-tracker", + "--summary", + summary_path, + ], + ) + assert result.exit_code == 0 + + # Read descriptors summary file + with open(summary_path, encoding="utf8") as file: + descriptors_summary = yaml.safe_load(file) + assert "emissions" not in descriptors_summary diff --git a/tests/test_eos_cli.py b/tests/test_eos_cli.py index 09510f8d..6e059100 100644 --- a/tests/test_eos_cli.py +++ b/tests/test_eos_cli.py @@ -300,3 +300,29 @@ def test_plot(tmp_path): ) assert result.exit_code == 0 assert plot_path.exists() + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + file_prefix = tmp_path / "NaCl" + summary_path = tmp_path / "NaCl-eos-summary.yml" + + result = runner.invoke( + app, + [ + "eos", + "--struct", + DATA_PATH / "NaCl.cif", + "--n-volumes", + 4, + "--no-tracker", + "--file-prefix", + file_prefix, + ], + ) + assert result.exit_code == 0 + + # Read eos summary file + with open(summary_path, encoding="utf8") as file: + eos_summary = yaml.safe_load(file) + assert "emissions" not in eos_summary diff --git a/tests/test_geomopt_cli.py b/tests/test_geomopt_cli.py index b020d409..9b25a265 100644 --- a/tests/test_geomopt_cli.py +++ b/tests/test_geomopt_cli.py @@ -709,3 +709,32 @@ def test_symmetrize(tmp_path): 90.0, ] assert results_2.cell.cellpar() == pytest.approx(expected) + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + results_path = tmp_path / "NaCl-results.extxyz" + log_path = tmp_path / "test.log" + summary_path = tmp_path / "summary.yml" + + result = runner.invoke( + app, + [ + "geomopt", + "--struct", + DATA_PATH / "NaCl.cif", + "--out", + results_path, + "--log", + log_path, + "--no-tracker", + "--summary", + summary_path, + ], + ) + assert result.exit_code == 0 + + # Read geomopt summary file + with open(summary_path, encoding="utf8") as file: + geomopt_summary = yaml.safe_load(file) + assert "emissions" not in geomopt_summary diff --git a/tests/test_md_cli.py b/tests/test_md_cli.py index 22f23750..5c6d441a 100644 --- a/tests/test_md_cli.py +++ b/tests/test_md_cli.py @@ -741,3 +741,32 @@ def test_auto_restart(tmp_path): # Includes header and steps 0, 3, and steps 5, 6, 7 assert len(lines) == 6 assert int(lines[-1].split()[0]) == 7 + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + file_prefix = tmp_path / "nvt" + summary_path = tmp_path / "nvt-md-summary.yml" + + result = runner.invoke( + app, + [ + "md", + "--ensemble", + "nvt", + "--struct", + DATA_PATH / "NaCl.cif", + "--no-tracker", + "--file-prefix", + file_prefix, + "--steps", + 1, + ], + ) + + assert result.exit_code == 0 + + # Read summary + with open(summary_path, encoding="utf8") as file: + summary = yaml.safe_load(file) + assert "emissions" not in summary diff --git a/tests/test_phonons_cli.py b/tests/test_phonons_cli.py index 9175469d..0eb64873 100644 --- a/tests/test_phonons_cli.py +++ b/tests/test_phonons_cli.py @@ -411,3 +411,28 @@ def test_invalid_traj_input(tmp_path): ) assert result.exit_code == 1 assert isinstance(result.exception, ValueError) + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + file_prefix = tmp_path / "NaCl" + summary_path = tmp_path / "NaCl-phonons-summary.yml" + + result = runner.invoke( + app, + [ + "phonons", + "--struct", + DATA_PATH / "NaCl.cif", + "--no-hdf5", + "--no-tracker", + "--file-prefix", + file_prefix, + ], + ) + assert result.exit_code == 0 + + # Read phonons summary file + with open(summary_path, encoding="utf8") as file: + phonon_summary = yaml.safe_load(file) + assert "emissions" not in phonon_summary diff --git a/tests/test_singlepoint_cli.py b/tests/test_singlepoint_cli.py index ed995158..b5f89958 100644 --- a/tests/test_singlepoint_cli.py +++ b/tests/test_singlepoint_cli.py @@ -391,3 +391,34 @@ def test_hessian(tmp_path): assert "mace_mp_hessian" in atoms.info assert "mace_stress" not in atoms.info assert atoms.info["mace_mp_hessian"].shape == (24, 8, 3) + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + results_path = tmp_path / "NaCl-results.extxyz" + log_path = tmp_path / "test.log" + summary_path = tmp_path / "summary.yml" + + result = runner.invoke( + app, + [ + "singlepoint", + "--struct", + DATA_PATH / "NaCl.cif", + "--properties", + "energy", + "--out", + results_path, + "--log", + log_path, + "--no-tracker", + "--summary", + summary_path, + ], + ) + assert result.exit_code == 0 + + # Read singlepoint summary file + with open(summary_path, encoding="utf8") as file: + sp_summary = yaml.safe_load(file) + assert "emissions" not in sp_summary diff --git a/tests/test_train_cli.py b/tests/test_train_cli.py index 75b94c2d..ddde4d20 100644 --- a/tests/test_train_cli.py +++ b/tests/test_train_cli.py @@ -260,3 +260,56 @@ def test_fine_tune_invalid_foundation(tmp_path): ) assert result.exit_code == 1 assert isinstance(result.exception, ValueError) + + +def test_no_carbon(tmp_path): + """Test disabling carbon tracking.""" + summary_path = tmp_path / "summary.yml" + model = "test.model" + compiled_model = "test_compiled.model" + results_path = "results" + checkpoints_path = "checkpoints" + logs_path = "logs" + log_path = Path("./train-log.yml").absolute() + summary_path = Path("./train-summary.yml").absolute() + + assert not Path(model).exists() + assert not Path(compiled_model).exists() + assert not Path(logs_path).exists() + assert not Path(results_path).exists() + assert not Path(checkpoints_path).exists() + assert not log_path.exists() + assert not summary_path.exists() + + config = write_tmp_config(DATA_PATH / "mlip_train.yml", tmp_path) + + result = runner.invoke( + app, + [ + "train", + "--mlip-config", + config, + "--no-tracker", + ], + ) + try: + assert result.exit_code == 0 + + assert summary_path.exists() + with open(summary_path, encoding="utf8") as file: + train_summary = yaml.safe_load(file) + assert "emissions" not in train_summary + finally: + # Tidy up models + Path(model).unlink(missing_ok=True) + Path(compiled_model).unlink(missing_ok=True) + + # Tidy up directories + shutil.rmtree(logs_path, ignore_errors=True) + shutil.rmtree(results_path, ignore_errors=True) + shutil.rmtree(checkpoints_path, ignore_errors=True) + + log_path.unlink(missing_ok=True) + summary_path.unlink(missing_ok=True) + + clear_log_handlers()