From 4df2c7539a56150b8264037a70257f0a04866954 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Thu, 12 Dec 2024 22:50:41 +0000
Subject: [PATCH 01/11] Save singlepoint units

---
 janus_core/calculations/base.py         | 30 +++++++++++++++++++++++++
 janus_core/calculations/single_point.py |  2 ++
 2 files changed, 32 insertions(+)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index eb8b3756..a13f86b0 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -18,6 +18,17 @@
 from janus_core.helpers.struct_io import input_structs
 from janus_core.helpers.utils import FileNameMixin, none_to_dict
 
+UNITS = {
+    "energy": "eV",
+    "forces": "ev/Ang",
+    "stress": "ev/Ang^3",
+    "time": "fs",
+    "real_time": "s",
+    "temperature": "K",
+    "pressure": "GPa",
+    "momenta": "(eV / u)^0.5",
+}
+
 
 class BaseCalculation(FileNameMixin):
     """
@@ -202,3 +213,22 @@ def __init__(
         self.tracker = config_tracker(
             self.logger, self.track_carbon, **self.tracker_kwargs
         )
+
+    def _set_units(self, keys: Sequence = ("energy", "forces", "stress")) -> None:
+        """
+        Save units to structure info.
+
+        Parameters
+        ----------
+        keys : Sequence
+            Keys for which to add units to structure info. Default is
+            ("energy", "forces", "stress").
+        """
+        if isinstance(self.struct, Sequence):
+            for image in self.struct:
+                for key in keys:
+                    image.info[f"{key}_units"] = UNITS[key]
+            return
+
+        for key in keys:
+            self.struct.info[f"{key}_units"] = UNITS[key]
diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py
index 953f0070..15e211ac 100644
--- a/janus_core/calculations/single_point.py
+++ b/janus_core/calculations/single_point.py
@@ -321,6 +321,8 @@ def run(self) -> CalcResults:
         if self.tracker:
             self.tracker.start_task("Single point")
 
+        self._set_units(self.properties)
+
         if "energy" in self.properties:
             self.results["energy"] = self._get_potential_energy()
         if "forces" in self.properties:

From 7c3bb42978f370322c98453a4ac698a6e9bd1bb0 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Thu, 12 Dec 2024 23:04:56 +0000
Subject: [PATCH 02/11] Save units for all calculations

---
 janus_core/calculations/eos.py      |  2 ++
 janus_core/calculations/geom_opt.py |  2 ++
 janus_core/calculations/md.py       | 17 +++++++++++++++--
 janus_core/calculations/phonons.py  |  2 ++
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/janus_core/calculations/eos.py b/janus_core/calculations/eos.py
index bd7bc863..5d23f759 100644
--- a/janus_core/calculations/eos.py
+++ b/janus_core/calculations/eos.py
@@ -290,6 +290,8 @@ def run(self) -> EoSResults:
             Dictionary containing equation of state ASE object, and the fitted minimum
             bulk modulus, volume, and energy.
         """
+        self._set_units()
+
         if self.minimize:
             if self.logger:
                 self.logger.info("Minimising initial structure")
diff --git a/janus_core/calculations/geom_opt.py b/janus_core/calculations/geom_opt.py
index 6077fddd..34b1e600 100644
--- a/janus_core/calculations/geom_opt.py
+++ b/janus_core/calculations/geom_opt.py
@@ -304,6 +304,8 @@ def run(self) -> None:
         if self.tracker:
             self.tracker.start_task("Geometry optimization")
 
+        self._set_units()
+
         converged = self.dyn.run(fmax=self.fmax, steps=self.steps)
 
         # Calculate current maximum force
diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
index 4a1119e0..2629b228 100644
--- a/janus_core/calculations/md.py
+++ b/janus_core/calculations/md.py
@@ -520,7 +520,7 @@ def _set_info(self) -> None:
         """Set time in fs, current dynamics step, and density to info."""
         time = (self.offset * self.timestep + self.dyn.get_time()) / units.fs
         step = self.offset + self.dyn.nsteps
-        self.dyn.atoms.info["time_fs"] = time
+        self.dyn.atoms.info["time"] = time
         self.dyn.atoms.info["step"] = step
         try:
             density = (
@@ -769,7 +769,7 @@ def get_stats(self) -> dict[str, float]:
         return {
             "Step": self.dyn.atoms.info["step"],
             "Real_Time": real_time.total_seconds(),
-            "Time": self.dyn.atoms.info["time_fs"],
+            "Time": self.dyn.atoms.info["time"],
             "Epot/N": e_pot,
             "EKin/N": e_kin,
             "T": current_temp,
@@ -1021,6 +1021,19 @@ def _write_restart(self) -> None:
 
     def run(self) -> None:
         """Run molecular dynamics simulation and/or temperature ramp."""
+        unit_keys = (
+            "energy",
+            "forces",
+            "stress",
+            "time",
+            "real_time",
+            "temperature",
+            "pressure",
+            "density",
+            "momenta",
+        )
+        self._set_units(unit_keys)
+
         if not self.restart:
             if self.minimize:
                 self._optimize_structure()
diff --git a/janus_core/calculations/phonons.py b/janus_core/calculations/phonons.py
index fd936fb6..6438fdd4 100644
--- a/janus_core/calculations/phonons.py
+++ b/janus_core/calculations/phonons.py
@@ -435,6 +435,8 @@ def calc_force_constants(
         if self.tracker:
             self.tracker.start_task("Phonon calculation")
 
+        self._set_units()
+
         cell = self._ASE_to_PhonopyAtoms(self.struct)
 
         if len(self.supercell) == 3:

From 0cae8d1bb764f1b7e5b4379b5a383893196d1d0a Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Fri, 3 Jan 2025 18:46:21 +0000
Subject: [PATCH 03/11] Add hessian units

---
 janus_core/calculations/base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index a13f86b0..13f72974 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -22,6 +22,7 @@
     "energy": "eV",
     "forces": "ev/Ang",
     "stress": "ev/Ang^3",
+    "hessian": "ev/Ang^2",
     "time": "fs",
     "real_time": "s",
     "temperature": "K",

From d7a92550ac627a3cca57c60973fbb95ea1b0d0ae Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Fri, 3 Jan 2025 18:46:46 +0000
Subject: [PATCH 04/11] Fix ASE unit definition

---
 janus_core/calculations/md.py        | 4 +++-
 janus_core/processing/observables.py | 5 ++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
index 2629b228..9b5a58df 100644
--- a/janus_core/calculations/md.py
+++ b/janus_core/calculations/md.py
@@ -12,7 +12,7 @@
 from typing import Any
 from warnings import warn
 
-from ase import Atoms, units
+from ase import Atoms
 from ase.geometry.analysis import Analysis
 from ase.io import read
 from ase.md.langevin import Langevin
@@ -23,6 +23,7 @@
     ZeroRotation,
 )
 from ase.md.verlet import VelocityVerlet
+from ase.units import create_units
 import numpy as np
 import yaml
 
@@ -43,6 +44,7 @@
 from janus_core.processing.correlator import Correlation
 from janus_core.processing.post_process import compute_rdf, compute_vaf
 
+units = create_units("2014")
 DENS_FACT = (units.m / 1.0e2) ** 3 / units.mol
 
 
diff --git a/janus_core/processing/observables.py b/janus_core/processing/observables.py
index 52fe0b91..38a946f9 100644
--- a/janus_core/processing/observables.py
+++ b/janus_core/processing/observables.py
@@ -5,13 +5,16 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
-from ase import Atoms, units
+from ase import Atoms
+from ase.units import create_units
 
 if TYPE_CHECKING:
     from janus_core.helpers.janus_types import SliceLike
 
 from janus_core.helpers.utils import slicelike_to_startstopstep
 
+units = create_units("2014")
+
 
 # pylint: disable=too-few-public-methods
 class Observable(ABC):

From ce1cbff30edb1c5041d04a10a0e815f22f7eb5d0 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Fri, 3 Jan 2025 19:12:48 +0000
Subject: [PATCH 05/11] Add density units

---
 janus_core/calculations/base.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index 13f72974..378913f8 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -27,7 +27,8 @@
     "real_time": "s",
     "temperature": "K",
     "pressure": "GPa",
-    "momenta": "(eV / u)^0.5",
+    "momenta": "(eV/u)^0.5",
+    "density": "g/cm^3",
 }
 
 

From 439c2658d9488ee48553800de3446f317cda08c0 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Mon, 6 Jan 2025 12:15:38 +0000
Subject: [PATCH 06/11] Save units as dict

---
 janus_core/calculations/base.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index 378913f8..b693a6b4 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -228,9 +228,6 @@ def _set_units(self, keys: Sequence = ("energy", "forces", "stress")) -> None:
         """
         if isinstance(self.struct, Sequence):
             for image in self.struct:
-                for key in keys:
-                    image.info[f"{key}_units"] = UNITS[key]
-            return
-
-        for key in keys:
-            self.struct.info[f"{key}_units"] = UNITS[key]
+                image.info["units"] = {key: UNITS[key] for key in keys}
+        else:
+            self.struct.info["units"] = {key: UNITS[key] for key in keys}

From a359fb620f80adc78c0e1a1c8e9e1e49504be3ce Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Mon, 6 Jan 2025 12:42:33 +0000
Subject: [PATCH 07/11] Test units

---
 tests/test_geomopt_cli.py     | 29 +++++++++++++++++++++++++++++
 tests/test_md_cli.py          | 17 +++++++++++++++++
 tests/test_singlepoint_cli.py |  6 ++++++
 3 files changed, 52 insertions(+)

diff --git a/tests/test_geomopt_cli.py b/tests/test_geomopt_cli.py
index dd2769f6..69d9af27 100644
--- a/tests/test_geomopt_cli.py
+++ b/tests/test_geomopt_cli.py
@@ -740,3 +740,32 @@ def test_no_carbon(tmp_path):
     with open(summary_path, encoding="utf8") as file:
         geomopt_summary = yaml.safe_load(file)
     assert "emissions" not in geomopt_summary
+
+
+def test_units(tmp_path):
+    """Test correct units are saved."""
+    results_path = tmp_path / "NaCl-opt.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,
+            "--summary",
+            summary_path,
+        ],
+    )
+    assert result.exit_code == 0
+
+    atoms = read(results_path)
+    expected_units = {"energy": "eV", "forces": "ev/Ang", "stress": "ev/Ang^3"}
+    assert "units" in atoms.info
+    for prop, units in expected_units.items():
+        assert atoms.info["units"][prop] == units
diff --git a/tests/test_md_cli.py b/tests/test_md_cli.py
index 1423c6b5..a7793a0b 100644
--- a/tests/test_md_cli.py
+++ b/tests/test_md_cli.py
@@ -108,6 +108,23 @@ def test_md(ensemble):
         assert "momenta" in atoms.arrays
         assert "masses" in atoms.arrays
 
+        expected_units = {
+            "time": "fs",
+            "real_time": "s",
+            "energy": "eV",
+            "forces": "ev/Ang",
+            "stress": "ev/Ang^3",
+            "temperature": "K",
+            "density": "g/cm^3",
+            "momenta": "(eV/u)^0.5",
+        }
+        if ensemble in ("nvt", "nvt-nh"):
+            expected_units["pressure"] = "GPa"
+
+        assert "units" in atoms.info
+        for prop, units in expected_units.items():
+            assert atoms.info["units"][prop] == units
+
     finally:
         final_path.unlink(missing_ok=True)
         restart_path.unlink(missing_ok=True)
diff --git a/tests/test_singlepoint_cli.py b/tests/test_singlepoint_cli.py
index 2476e5a7..099fe6b7 100644
--- a/tests/test_singlepoint_cli.py
+++ b/tests/test_singlepoint_cli.py
@@ -72,6 +72,11 @@ def test_singlepoint():
         assert "system_name" in atoms.info
         assert atoms.info["system_name"] == "NaCl"
 
+        expected_units = {"energy": "eV", "forces": "ev/Ang", "stress": "ev/Ang^3"}
+        assert "units" in atoms.info
+        for prop, units in expected_units.items():
+            assert atoms.info["units"][prop] == units
+
         clear_log_handlers()
 
 
@@ -399,6 +404,7 @@ 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)
+    assert atoms.info["units"]["hessian"] == "ev/Ang^2"
 
 
 def test_no_carbon(tmp_path):

From 4a21da76cc48310bb422a9f23487c062ed005470 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Wed, 8 Jan 2025 16:34:23 +0000
Subject: [PATCH 08/11] Fix momentum units

---
 janus_core/calculations/base.py | 2 +-
 tests/test_md_cli.py            | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index b693a6b4..64952902 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -27,7 +27,7 @@
     "real_time": "s",
     "temperature": "K",
     "pressure": "GPa",
-    "momenta": "(eV/u)^0.5",
+    "momenta": "(eV*u)^0.5",
     "density": "g/cm^3",
 }
 
diff --git a/tests/test_md_cli.py b/tests/test_md_cli.py
index a7793a0b..771d2094 100644
--- a/tests/test_md_cli.py
+++ b/tests/test_md_cli.py
@@ -116,7 +116,7 @@ def test_md(ensemble):
             "stress": "ev/Ang^3",
             "temperature": "K",
             "density": "g/cm^3",
-            "momenta": "(eV/u)^0.5",
+            "momenta": "(eV*u)^0.5",
         }
         if ensemble in ("nvt", "nvt-nh"):
             expected_units["pressure"] = "GPa"

From 1fa9bc4e0cc82c56dbbd27679802a325c5c9b249 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Mon, 13 Jan 2025 17:34:44 +0000
Subject: [PATCH 09/11] Update janus_core/calculations/base.py

Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com>
---
 janus_core/calculations/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index 64952902..77ce4aa2 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -216,7 +216,7 @@ def __init__(
             self.logger, self.track_carbon, **self.tracker_kwargs
         )
 
-    def _set_units(self, keys: Sequence = ("energy", "forces", "stress")) -> None:
+    def _set_units(self, keys: Sequence[str] = ("energy", "forces", "stress")) -> None:
         """
         Save units to structure info.
 

From 154223fc5ccf166700a1436ac409748f6f7e74d9 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Tue, 14 Jan 2025 11:36:44 +0000
Subject: [PATCH 10/11] Use Base units for MD stats

---
 janus_core/calculations/base.py |  1 +
 janus_core/calculations/md.py   | 41 ++++++++++++++++++---------------
 tests/test_md_cli.py            |  2 +-
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index 77ce4aa2..065662b2 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -29,6 +29,7 @@
     "pressure": "GPa",
     "momenta": "(eV*u)^0.5",
     "density": "g/cm^3",
+    "volume": "Ang^3",
 }
 
 
diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
index 9b5a58df..6fe68c96 100644
--- a/janus_core/calculations/md.py
+++ b/janus_core/calculations/md.py
@@ -27,7 +27,7 @@
 import numpy as np
 import yaml
 
-from janus_core.calculations.base import BaseCalculation
+from janus_core.calculations.base import UNITS, BaseCalculation
 from janus_core.calculations.geom_opt import GeomOpt
 from janus_core.helpers.janus_types import (
     Architectures,
@@ -799,21 +799,21 @@ def unit_info(self) -> dict[str, str]:
         """
         return {
             "Step": None,
-            "Real_Time": "s",
-            "Time": "fs",
-            "Epot/N": "eV",
-            "EKin/N": "eV",
-            "T": "K",
-            "ETot/N": "eV",
-            "Density": "g/cm^3",
-            "Volume": "A^3",
-            "P": "GPa",
-            "Pxx": "GPa",
-            "Pyy": "GPa",
-            "Pzz": "GPa",
-            "Pyz": "GPa",
-            "Pxz": "GPa",
-            "Pxy": "GPa",
+            "Real_Time": UNITS["real_time"],
+            "Time": UNITS["time"],
+            "Epot/N": UNITS["energy"],
+            "EKin/N": UNITS["energy"],
+            "T": UNITS["temperature"],
+            "ETot/N": UNITS["energy"],
+            "Density": UNITS["density"],
+            "Volume": UNITS["volume"],
+            "P": UNITS["pressure"],
+            "Pxx": UNITS["pressure"],
+            "Pyy": UNITS["pressure"],
+            "Pzz": UNITS["pressure"],
+            "Pyz": UNITS["pressure"],
+            "Pxz": UNITS["pressure"],
+            "Pxy": UNITS["pressure"],
         }
 
     @property
@@ -1277,7 +1277,10 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_P": "GPa", "Target_T": "K"}
+        return super().unit_info | {
+            "Target_P": UNITS["pressure"],
+            "Target_T": UNITS["temperature"],
+        }
 
     @property
     def default_formats(self) -> dict[str, str]:
@@ -1374,7 +1377,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_T": "K"}
+        return super().unit_info | {"Target_T": UNITS["temperature"]}
 
     @property
     def default_formats(self) -> dict[str, str]:
@@ -1517,7 +1520,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_T": "K"}
+        return super().unit_info | {"Target_T": UNITS["temperature"]}
 
     @property
     def default_formats(self) -> dict[str, str]:
diff --git a/tests/test_md_cli.py b/tests/test_md_cli.py
index 771d2094..1161ad46 100644
--- a/tests/test_md_cli.py
+++ b/tests/test_md_cli.py
@@ -170,7 +170,7 @@ def test_log(tmp_path):
         assert len(lines) == 22
 
         # Test constant volume
-        assert lines[0].split(" | ")[8] == "Volume [A^3]"
+        assert lines[0].split(" | ")[8] == "Volume [Ang^3]"
         init_volume = float(lines[1].split()[8])
         final_volume = float(lines[-1].split()[8])
         assert init_volume == 179.406144

From 48e014e95e0f3d836e655fd0ebc9ece25f42a2ba Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Tue, 14 Jan 2025 14:38:52 +0000
Subject: [PATCH 11/11] Rename units for clarity

---
 janus_core/calculations/base.py         |  4 ++-
 janus_core/calculations/eos.py          |  2 +-
 janus_core/calculations/geom_opt.py     |  2 +-
 janus_core/calculations/md.py           | 43 +++++++++++++------------
 janus_core/calculations/phonons.py      |  2 +-
 janus_core/calculations/single_point.py |  2 +-
 6 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
index 065662b2..554345be 100644
--- a/janus_core/calculations/base.py
+++ b/janus_core/calculations/base.py
@@ -217,7 +217,9 @@ def __init__(
             self.logger, self.track_carbon, **self.tracker_kwargs
         )
 
-    def _set_units(self, keys: Sequence[str] = ("energy", "forces", "stress")) -> None:
+    def _set_info_units(
+        self, keys: Sequence[str] = ("energy", "forces", "stress")
+    ) -> None:
         """
         Save units to structure info.
 
diff --git a/janus_core/calculations/eos.py b/janus_core/calculations/eos.py
index 5d23f759..e6bddf6a 100644
--- a/janus_core/calculations/eos.py
+++ b/janus_core/calculations/eos.py
@@ -290,7 +290,7 @@ def run(self) -> EoSResults:
             Dictionary containing equation of state ASE object, and the fitted minimum
             bulk modulus, volume, and energy.
         """
-        self._set_units()
+        self._set_info_units()
 
         if self.minimize:
             if self.logger:
diff --git a/janus_core/calculations/geom_opt.py b/janus_core/calculations/geom_opt.py
index 34b1e600..f4af4eec 100644
--- a/janus_core/calculations/geom_opt.py
+++ b/janus_core/calculations/geom_opt.py
@@ -304,7 +304,7 @@ def run(self) -> None:
         if self.tracker:
             self.tracker.start_task("Geometry optimization")
 
-        self._set_units()
+        self._set_info_units()
 
         converged = self.dyn.run(fmax=self.fmax, steps=self.steps)
 
diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
index 6fe68c96..cf71924a 100644
--- a/janus_core/calculations/md.py
+++ b/janus_core/calculations/md.py
@@ -27,7 +27,8 @@
 import numpy as np
 import yaml
 
-from janus_core.calculations.base import UNITS, BaseCalculation
+from janus_core.calculations.base import UNITS as JANUS_UNITS
+from janus_core.calculations.base import BaseCalculation
 from janus_core.calculations.geom_opt import GeomOpt
 from janus_core.helpers.janus_types import (
     Architectures,
@@ -799,21 +800,21 @@ def unit_info(self) -> dict[str, str]:
         """
         return {
             "Step": None,
-            "Real_Time": UNITS["real_time"],
-            "Time": UNITS["time"],
-            "Epot/N": UNITS["energy"],
-            "EKin/N": UNITS["energy"],
-            "T": UNITS["temperature"],
-            "ETot/N": UNITS["energy"],
-            "Density": UNITS["density"],
-            "Volume": UNITS["volume"],
-            "P": UNITS["pressure"],
-            "Pxx": UNITS["pressure"],
-            "Pyy": UNITS["pressure"],
-            "Pzz": UNITS["pressure"],
-            "Pyz": UNITS["pressure"],
-            "Pxz": UNITS["pressure"],
-            "Pxy": UNITS["pressure"],
+            "Real_Time": JANUS_UNITS["real_time"],
+            "Time": JANUS_UNITS["time"],
+            "Epot/N": JANUS_UNITS["energy"],
+            "EKin/N": JANUS_UNITS["energy"],
+            "T": JANUS_UNITS["temperature"],
+            "ETot/N": JANUS_UNITS["energy"],
+            "Density": JANUS_UNITS["density"],
+            "Volume": JANUS_UNITS["volume"],
+            "P": JANUS_UNITS["pressure"],
+            "Pxx": JANUS_UNITS["pressure"],
+            "Pyy": JANUS_UNITS["pressure"],
+            "Pzz": JANUS_UNITS["pressure"],
+            "Pyz": JANUS_UNITS["pressure"],
+            "Pxz": JANUS_UNITS["pressure"],
+            "Pxy": JANUS_UNITS["pressure"],
         }
 
     @property
@@ -1034,7 +1035,7 @@ def run(self) -> None:
             "density",
             "momenta",
         )
-        self._set_units(unit_keys)
+        self._set_info_units(unit_keys)
 
         if not self.restart:
             if self.minimize:
@@ -1278,8 +1279,8 @@ def unit_info(self) -> dict[str, str]:
             Units attached to statistical properties.
         """
         return super().unit_info | {
-            "Target_P": UNITS["pressure"],
-            "Target_T": UNITS["temperature"],
+            "Target_P": JANUS_UNITS["pressure"],
+            "Target_T": JANUS_UNITS["temperature"],
         }
 
     @property
@@ -1377,7 +1378,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_T": UNITS["temperature"]}
+        return super().unit_info | {"Target_T": JANUS_UNITS["temperature"]}
 
     @property
     def default_formats(self) -> dict[str, str]:
@@ -1520,7 +1521,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_T": UNITS["temperature"]}
+        return super().unit_info | {"Target_T": JANUS_UNITS["temperature"]}
 
     @property
     def default_formats(self) -> dict[str, str]:
diff --git a/janus_core/calculations/phonons.py b/janus_core/calculations/phonons.py
index 6438fdd4..7f3e6062 100644
--- a/janus_core/calculations/phonons.py
+++ b/janus_core/calculations/phonons.py
@@ -435,7 +435,7 @@ def calc_force_constants(
         if self.tracker:
             self.tracker.start_task("Phonon calculation")
 
-        self._set_units()
+        self._set_info_units()
 
         cell = self._ASE_to_PhonopyAtoms(self.struct)
 
diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py
index 15e211ac..806bf622 100644
--- a/janus_core/calculations/single_point.py
+++ b/janus_core/calculations/single_point.py
@@ -321,7 +321,7 @@ def run(self) -> CalcResults:
         if self.tracker:
             self.tracker.start_task("Single point")
 
-        self._set_units(self.properties)
+        self._set_info_units(self.properties)
 
         if "energy" in self.properties:
             self.results["energy"] = self._get_potential_energy()