Skip to content

Commit

Permalink
Make carbon tracking optional
Browse files Browse the repository at this point in the history
  • Loading branch information
ElliottKasoar committed Oct 29, 2024
1 parent 4781620 commit aaf0536
Show file tree
Hide file tree
Showing 23 changed files with 377 additions and 21 deletions.
10 changes: 9 additions & 1 deletion janus_core/calculations/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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]
Expand All @@ -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:
Expand Down Expand Up @@ -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
)
10 changes: 9 additions & 1 deletion janus_core/calculations/descriptors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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):
Expand All @@ -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,
Expand Down
14 changes: 12 additions & 2 deletions janus_core/calculations/eos.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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:
Expand All @@ -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()
Expand Down Expand Up @@ -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")
10 changes: 9 additions & 1 deletion janus_core/calculations/geom_opt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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")
12 changes: 11 additions & 1 deletion janus_core/calculations/md.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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):
Expand Down
Loading

0 comments on commit aaf0536

Please sign in to comment.