diff --git a/a3fe/__init__.py b/a3fe/__init__.py
index 52fa4d6..0b3d19a 100644
--- a/a3fe/__init__.py
+++ b/a3fe/__init__.py
@@ -24,10 +24,11 @@
     Simulation,
     Stage,
     StageType,
-    SystemPreparationConfig,
     enums,
 )
 
+from .configuration import SystemPreparationConfig
+
 _sys.modules["EnsEquil"] = _sys.modules["a3fe"]
 
 # A3FE can open many files due to the use of multiprocessing and
diff --git a/a3fe/configuration/__init__.py b/a3fe/configuration/__init__.py
new file mode 100644
index 0000000..b8030db
--- /dev/null
+++ b/a3fe/configuration/__init__.py
@@ -0,0 +1,3 @@
+"""Pydantic configuration classes for the a3fe package."""
+
+from .system_prep_config import SystemPreparationConfig
diff --git a/a3fe/configuration/system_prep_config.py b/a3fe/configuration/system_prep_config.py
new file mode 100644
index 0000000..f91e879
--- /dev/null
+++ b/a3fe/configuration/system_prep_config.py
@@ -0,0 +1,251 @@
+"""
+Configuration classes for system preparation.
+"""
+
+__all__ = [
+    "SystemPreparationConfig",
+]
+
+import yaml as _yaml
+
+from pydantic import BaseModel as _BaseModel
+from pydantic import Field as _Field
+from pydantic import ConfigDict as _ConfigDict
+
+from ..run.enums import StageType as _StageType
+from ..run.enums import LegType as _LegType
+
+
+class SystemPreparationConfig(_BaseModel):
+    """
+    Pydantic model for holding system preparation configuration.
+
+    Attributes
+    ----------
+    slurm: bool
+        Whether to use SLURM for the preparation.
+    forcefields : dict
+        Forcefields to use for the ligand, protein, and water.
+    water_model : str
+        Water model to use.
+    ion_conc : float
+        Ion concentration in M.
+    steps : int
+        Number of steps for the minimisation.
+    runtime_short_nvt : int
+        Runtime for the short NVT equilibration in ps.
+    runtime_nvt : int
+        Runtime for the NVT equilibration in ps.
+    end_temp : float
+        End temperature for the NVT equilibration in K.
+    runtime_npt : int
+        Runtime for the NPT equilibration in ps.
+    runtime_npt_unrestrained : int
+        Runtime for the unrestrained NPT equilibration in ps.
+    ensemble_equilibration_time : int
+        Ensemble equilibration time in ps.
+    append_to_ligand_selection: str
+        If this is a bound leg, this appends the supplied string to the default atom
+        selection which chooses the atoms in the ligand to consider as potential anchor
+        points. The default atom selection is f'resname {ligand_resname} and not name H*'.
+        Uses the mdanalysis atom selection language. For example, 'not name O*' will result
+        in an atom selection of f'resname {ligand_resname} and not name H* and not name O*'.
+    use_same_restraints: bool
+        If True, the same restraints will be used for all of the bound leg repeats - by default
+        , the restraints generated for the first repeat are used. This allows meaningful
+        comparison between repeats for the bound leg. If False, the unique restraints are
+        generated for each repeat.
+    """
+
+    slurm: bool = _Field(True)
+    forcefields: dict = {
+        "ligand": "openff_unconstrained-2.0.0",
+        "protein": "ff14SB",
+        "water": "tip3p",
+    }
+    water_model: str = "tip3p"
+    ion_conc: float = _Field(0.15, ge=0, lt=1)  # M
+    steps: int = _Field(1000, gt=0, lt=100_000)  # This is the default for _BSS
+    runtime_short_nvt: int = _Field(5, gt=0, lt=500)  # ps
+    runtime_nvt: int = _Field(50, gt=0, lt=5_000)  # ps
+    end_temp: float = _Field(298.15, gt=0, lt=350)  # K
+    runtime_npt: int = _Field(400, gt=0, lt=40_000)  # ps
+    runtime_npt_unrestrained: int = _Field(1000, gt=0, lt=100_000)  # ps
+    ensemble_equilibration_time: int = _Field(5000, gt=0, lt=50_000)  # ps
+    append_to_ligand_selection: str = _Field(
+        "",
+        description="Atom selection to append to the ligand selection during restraint searching.",
+    )
+    use_same_restraints: bool = _Field(
+        True,
+        description="Whether to use the same restraints for all repeats of the bound leg. Note "
+        "that this should be used if you plan to run adaptively.",
+    )
+    lambda_values: dict = {
+        _LegType.BOUND: {
+            _StageType.RESTRAIN: [0.000, 0.125, 0.250, 0.375, 0.500, 1.000],
+            _StageType.DISCHARGE: [
+                0.000,
+                0.143,
+                0.286,
+                0.429,
+                0.571,
+                0.714,
+                0.857,
+                1.000,
+            ],
+            _StageType.VANISH: [
+                0.000,
+                0.025,
+                0.050,
+                0.075,
+                0.100,
+                0.125,
+                0.150,
+                0.175,
+                0.200,
+                0.225,
+                0.250,
+                0.275,
+                0.300,
+                0.325,
+                0.350,
+                0.375,
+                0.400,
+                0.425,
+                0.450,
+                0.475,
+                0.500,
+                0.525,
+                0.550,
+                0.575,
+                0.600,
+                0.625,
+                0.650,
+                0.675,
+                0.700,
+                0.725,
+                0.750,
+                0.800,
+                0.850,
+                0.900,
+                0.950,
+                1.000,
+            ],
+        },
+        _LegType.FREE: {
+            _StageType.DISCHARGE: [
+                0.000,
+                0.143,
+                0.286,
+                0.429,
+                0.571,
+                0.714,
+                0.857,
+                1.000,
+            ],
+            _StageType.VANISH: [
+                0.000,
+                0.028,
+                0.056,
+                0.111,
+                0.167,
+                0.222,
+                0.278,
+                0.333,
+                0.389,
+                0.444,
+                0.500,
+                0.556,
+                0.611,
+                0.667,
+                0.722,
+                0.778,
+                0.889,
+                1.000,
+            ],
+        },
+    }
+
+    model_config = _ConfigDict(extra="forbid", validate_assignment=True)
+
+    def get_tot_simtime(self, n_runs: int, leg_type: _LegType) -> int:
+        """
+        Get the total simulation time for the ensemble equilibration.
+
+        Parameters
+        ----------
+        n_runs : int
+            Number of ensemble equilibration runs.
+        leg_type : LegType
+            The type of the leg.
+
+        Returns
+        -------
+        int
+            Total simulation time in ps.
+        """
+
+        # See functions below for where these times are used.
+        tot_simtime = 0
+        tot_simtime += self.runtime_short_nvt
+        tot_simtime += (
+            self.runtime_nvt * 2 if leg_type == _LegType.BOUND else self.runtime_nvt
+        )
+        tot_simtime += self.runtime_npt * 2
+        tot_simtime += self.runtime_npt_unrestrained
+        tot_simtime += self.ensemble_equilibration_time * n_runs
+        return tot_simtime
+
+    def dump(self, save_dir: str, leg_type: _LegType) -> None:
+        """
+        Save the configuration to a YAML file.
+
+        Parameters
+        ----------
+        save_dir : str
+            Directory to save the YAML file to.
+
+        leg_type : LegType
+            The type of the leg. Used to name the YAML file.
+        """
+        # First, convert to dict
+        model_dict = self.model_dump()
+
+        # Save the dict to a YAML file
+        save_path = save_dir + "/" + self.get_file_name(leg_type)
+        with open(save_path, "w") as f:
+            _yaml.dump(model_dict, f, default_flow_style=False)
+
+    @classmethod
+    def load(cls, save_dir: str, leg_type: _LegType) -> "SystemPreparationConfig":
+        """
+        Load the configuration from a YAML file.
+
+        Parameters
+        ----------
+        save_dir : str
+            Directory to load the YAML file from.
+
+        leg_type : LegType
+            The type of the leg. Used to decide the name of the YAML
+            file to load.
+
+        Returns
+        -------
+        SystemPreparationConfig
+            Loaded configuration.
+        """
+
+        # Load the dict from the YAML file
+        load_path = save_dir + "/" + cls.get_file_name(leg_type)
+        with open(load_path, "r") as f:
+            model_dict = _yaml.load(f, Loader=_yaml.FullLoader)
+
+        # Create the model from the dict
+        return cls.model_validate(model_dict)
+
+    @staticmethod
+    def get_file_name(leg_type: _LegType) -> str:
+        """Get the name of the YAML file for the configuration."""
+        return f"system_preparation_config_{leg_type.name.lower()}.yaml"
diff --git a/a3fe/run/__init__.py b/a3fe/run/__init__.py
index 0df5aa4..7236579 100644
--- a/a3fe/run/__init__.py
+++ b/a3fe/run/__init__.py
@@ -5,4 +5,4 @@
 from .leg import Leg
 from .simulation import Simulation
 from .stage import Stage
-from .system_prep import SystemPreparationConfig
+from ..configuration import SystemPreparationConfig
diff --git a/a3fe/run/calc_set.py b/a3fe/run/calc_set.py
index 73126e1..b0d9901 100644
--- a/a3fe/run/calc_set.py
+++ b/a3fe/run/calc_set.py
@@ -18,7 +18,7 @@
 from ._simulation_runner import SimulationRunner as _SimulationRunner
 from ._utils import SimulationRunnerIterator as _SimulationRunnerIterator
 from .calculation import Calculation as _Calculation
-from .system_prep import SystemPreparationConfig as _SystemPreparationConfig
+from ..configuration import SystemPreparationConfig as _SystemPreparationConfig
 
 
 class CalcSet(_SimulationRunner):
diff --git a/a3fe/run/calculation.py b/a3fe/run/calculation.py
index 5755e3d..9561ec6 100644
--- a/a3fe/run/calculation.py
+++ b/a3fe/run/calculation.py
@@ -13,7 +13,7 @@
 from .enums import LegType as _LegType
 from .enums import PreparationStage as _PreparationStage
 from .leg import Leg as _Leg
-from .system_prep import SystemPreparationConfig as _SystemPreparationConfig
+from ..configuration import SystemPreparationConfig as _SystemPreparationConfig
 
 
 class Calculation(_SimulationRunner):
diff --git a/a3fe/run/enums.py b/a3fe/run/enums.py
index 6d30797..fde628f 100644
--- a/a3fe/run/enums.py
+++ b/a3fe/run/enums.py
@@ -2,6 +2,9 @@
 
 from enum import Enum as _Enum
 from typing import List as _List
+import yaml as _yaml
+
+from typing import Any as _Any
 
 __all__ = [
     "JobStatus",
@@ -11,7 +14,45 @@
 ]
 
 
-class JobStatus(_Enum):
+class _YamlSerialisableEnum(_Enum):
+    """A base class for enums that can be serialised to and deserialised from YAML."""
+    @classmethod
+    def to_yaml(cls, dumper: _yaml.Dumper, data: _Any) -> _yaml.nodes.ScalarNode:
+        return dumper.represent_scalar(
+            f"!{cls.__name__}", f"{cls.__name__}.{data.name}"
+        )
+
+    @classmethod
+    def from_yaml(cls, loader: _yaml.Loader, node: _yaml.nodes.ScalarNode) -> _Any:
+        value = loader.construct_scalar(node)
+        enum_name, member_name = value.split(".")
+        enum_class = globals()[enum_name]
+        return enum_class[member_name]
+
+
+# Register the custom representers and constructors for _YamlSerialisableEnum
+def _yaml_enum_representer(
+    dumper: _yaml.Dumper, data: _YamlSerialisableEnum
+) -> _yaml.nodes.ScalarNode:
+    return dumper.represent_scalar(
+        f"!{data.__class__.__name__}", f"{data.__class__.__name__}.{data.name}"
+    )
+
+
+def _yaml_enum_constructor(
+    loader: _yaml.Loader, suffix: str, node: _yaml.nodes.ScalarNode
+) -> _Any:
+    value = loader.construct_scalar(node)
+    enum_name, member_name = value.split(".")
+    enum_class = globals()[enum_name]
+    return enum_class[member_name]
+
+
+_yaml.add_multi_representer(_YamlSerialisableEnum, _yaml_enum_representer)
+_yaml.add_multi_constructor("!", _yaml_enum_constructor)
+
+
+class JobStatus(_YamlSerialisableEnum):
     """An enumeration of the possible job statuses"""
 
     NONE = 0
@@ -21,7 +62,7 @@ class JobStatus(_Enum):
     KILLED = 4
 
 
-class StageType(_Enum):
+class StageType(_YamlSerialisableEnum):
     """Enumeration of the types of stage."""
 
     RESTRAIN = 1
@@ -41,14 +82,14 @@ def bss_perturbation_type(self) -> str:
             raise ValueError("Unknown stage type.")
 
 
-class LegType(_Enum):
+class LegType(_YamlSerialisableEnum):
     """The type of leg in the calculation."""
 
     BOUND = 1
     FREE = 2
 
 
-class PreparationStage(_Enum):
+class PreparationStage(_YamlSerialisableEnum):
     """The stage of preparation of the input files."""
 
     STRUCTURES_ONLY = 1
diff --git a/a3fe/run/leg.py b/a3fe/run/leg.py
index 401a152..e5c8e2a 100644
--- a/a3fe/run/leg.py
+++ b/a3fe/run/leg.py
@@ -35,7 +35,7 @@
 from .enums import PreparationStage as _PreparationStage
 from .enums import StageType as _StageType
 from .stage import Stage as _Stage
-from .system_prep import SystemPreparationConfig as _SystemPreparationConfig
+from ..configuration import SystemPreparationConfig as _SystemPreparationConfig
 
 
 class Leg(_SimulationRunner):
@@ -215,7 +215,7 @@ def setup(
         cfg = (
             sysprep_config if sysprep_config is not None else _SystemPreparationConfig()
         )
-        cfg.save_pickle(self.input_dir, self.leg_type)
+        cfg.dump(self.input_dir, self.leg_type)
 
         # Create input directories, parameterise, solvate, minimise, heat and preequil, all
         # depending on the input files present.
@@ -597,7 +597,7 @@ def run_ensemble_equilibration(
                 _subprocess.run(["cp", "-r", input_file, outdir], check=True)
 
             # Also write a pickle of the config to the output directory
-            sysprep_config.save_pickle(outdir, self.leg_type)
+            sysprep_config.dump(outdir, self.leg_type)
 
         if sysprep_config.slurm:
             if self.leg_type == _LegType.BOUND:
diff --git a/a3fe/run/stage.py b/a3fe/run/stage.py
index fa27a45..0820f82 100644
--- a/a3fe/run/stage.py
+++ b/a3fe/run/stage.py
@@ -764,7 +764,9 @@ def analyse(
         for win in self.lam_windows:
             if not win.equilibrated:
                 raise RuntimeError(
-                    "Not all lambda windows have equilibrated. Analysis cannot be performed."
+                    "Not all lambda windows have equilibrated. Analysis cannot be performed. "
+                    "If you are running non-adaptively, please use the `set_equilibration_time` method "
+                    "to set the equilibration time manually."
                 )
             if win.equil_time is None:
                 raise RuntimeError(
diff --git a/a3fe/run/system_prep.py b/a3fe/run/system_prep.py
index 3b6e6b6..19639e1 100644
--- a/a3fe/run/system_prep.py
+++ b/a3fe/run/system_prep.py
@@ -1,7 +1,6 @@
 """Functionality for running preparation simulations."""
 
 __all__ = [
-    "SystemPreparationConfig",
     "parameterise_input",
     "solvate_input",
     "minimise_input",
@@ -10,262 +9,19 @@
 ]
 
 import pathlib as _pathlib
-import pickle as _pkl
 import warnings as _warnings
 from typing import Optional as _Optional
 
 import BioSimSpace.Sandpit.Exscientia as _BSS
-from pydantic import BaseModel as _BaseModel
-from pydantic import Field as _Field
 
 from ..read._process_bss_systems import rename_lig as _rename_lig
 from ._utils import check_has_wat_and_box as _check_has_wat_and_box
 from .enums import LegType as _LegType
 from .enums import PreparationStage as _PreparationStage
-from .enums import StageType as _StageType
 
-
-class SystemPreparationConfig(_BaseModel):
-    """
-    Pydantic model for holding system preparation configuration.
-
-    Attributes
-    ----------
-    slurm: bool
-        Whether to use SLURM for the preparation.
-    forcefields : dict
-        Forcefields to use for the ligand, protein, and water.
-    water_model : str
-        Water model to use.
-    ion_conc : float
-        Ion concentration in M.
-    steps : int
-        Number of steps for the minimisation.
-    runtime_short_nvt : int
-        Runtime for the short NVT equilibration in ps.
-    runtime_nvt : int
-        Runtime for the NVT equilibration in ps.
-    end_temp : float
-        End temperature for the NVT equilibration in K.
-    runtime_npt : int
-        Runtime for the NPT equilibration in ps.
-    runtime_npt_unrestrained : int
-        Runtime for the unrestrained NPT equilibration in ps.
-    ensemble_equilibration_time : int
-        Ensemble equilibration time in ps.
-    append_to_ligand_selection: str
-        If this is a bound leg, this appends the supplied string to the default atom
-        selection which chooses the atoms in the ligand to consider as potential anchor
-        points. The default atom selection is f'resname {ligand_resname} and not name H*'.
-        Uses the mdanalysis atom selection language. For example, 'not name O*' will result
-        in an atom selection of f'resname {ligand_resname} and not name H* and not name O*'.
-    use_same_restraints: bool
-        If True, the same restraints will be used for all of the bound leg repeats - by default
-        , the restraints generated for the first repeat are used. This allows meaningful
-        comparison between repeats for the bound leg. If False, the unique restraints are
-        generated for each repeat.
-    """
-
-    slurm: bool = _Field(True)
-    forcefields: dict = {
-        "ligand": "openff_unconstrained-2.0.0",
-        "protein": "ff14SB",
-        "water": "tip3p",
-    }
-    water_model: str = "tip3p"
-    ion_conc: float = _Field(0.15, ge=0, lt=1)  # M
-    steps: int = _Field(1000, gt=0, lt=100_000)  # This is the default for _BSS
-    runtime_short_nvt: int = _Field(5, gt=0, lt=500)  # ps
-    runtime_nvt: int = _Field(50, gt=0, lt=5_000)  # ps
-    end_temp: float = _Field(298.15, gt=0, lt=350)  # K
-    runtime_npt: int = _Field(400, gt=0, lt=40_000)  # ps
-    runtime_npt_unrestrained: int = _Field(1000, gt=0, lt=100_000)  # ps
-    ensemble_equilibration_time: int = _Field(5000, gt=0, lt=50_000)  # ps
-    append_to_ligand_selection: str = _Field(
-        "",
-        description="Atom selection to append to the ligand selection during restraint searching.",
-    )
-    use_same_restraints: bool = _Field(
-        True,
-        description="Whether to use the same restraints for all repeats of the bound leg. Note "
-        "that this should be used if you plan to run adaptively.",
-    )
-    lambda_values: dict = {
-        _LegType.BOUND: {
-            _StageType.RESTRAIN: [0.000, 0.125, 0.250, 0.375, 0.500, 1.000],
-            _StageType.DISCHARGE: [
-                0.000,
-                0.143,
-                0.286,
-                0.429,
-                0.571,
-                0.714,
-                0.857,
-                1.000,
-            ],
-            _StageType.VANISH: [
-                0.000,
-                0.025,
-                0.050,
-                0.075,
-                0.100,
-                0.125,
-                0.150,
-                0.175,
-                0.200,
-                0.225,
-                0.250,
-                0.275,
-                0.300,
-                0.325,
-                0.350,
-                0.375,
-                0.400,
-                0.425,
-                0.450,
-                0.475,
-                0.500,
-                0.525,
-                0.550,
-                0.575,
-                0.600,
-                0.625,
-                0.650,
-                0.675,
-                0.700,
-                0.725,
-                0.750,
-                0.800,
-                0.850,
-                0.900,
-                0.950,
-                1.000,
-            ],
-        },
-        _LegType.FREE: {
-            _StageType.DISCHARGE: [
-                0.000,
-                0.143,
-                0.286,
-                0.429,
-                0.571,
-                0.714,
-                0.857,
-                1.000,
-            ],
-            _StageType.VANISH: [
-                0.000,
-                0.028,
-                0.056,
-                0.111,
-                0.167,
-                0.222,
-                0.278,
-                0.333,
-                0.389,
-                0.444,
-                0.500,
-                0.556,
-                0.611,
-                0.667,
-                0.722,
-                0.778,
-                0.889,
-                1.000,
-            ],
-        },
-    }
-
-    class Config:
-        """
-        Pydantic model configuration.
-        """
-
-        extra = "forbid"
-        validate_assignment = True
-
-    def get_tot_simtime(self, n_runs: int, leg_type: _LegType) -> int:
-        """
-        Get the total simulation time for the ensemble equilibration.
-
-        Parameters
-        ----------
-        n_runs : int
-            Number of ensemble equilibration runs.
-        leg_type : LegType
-            The type of the leg.
-
-        Returns
-        -------
-        int
-            Total simulation time in ps.
-        """
-
-        # See functions below for where these times are used.
-        tot_simtime = 0
-        tot_simtime += self.runtime_short_nvt
-        tot_simtime += (
-            self.runtime_nvt * 2 if leg_type == _LegType.BOUND else self.runtime_nvt
-        )
-        tot_simtime += self.runtime_npt * 2
-        tot_simtime += self.runtime_npt_unrestrained
-        tot_simtime += self.ensemble_equilibration_time * n_runs
-        return tot_simtime
-
-    def save_pickle(self, save_dir: str, leg_type: _LegType) -> None:
-        """
-        Save the configuration to a pickle file.
-
-        Parameters
-        ----------
-        save_dir : str
-            Directory to save the pickle file to.
-
-        leg_type : LegType
-            The type of the leg. Used to name the pickle file.
-        """
-        # First, convert to dict
-        model_dict = self.model_dump()
-
-        # Save the dict to a pickle file
-        save_path = save_dir + "/" + self.get_file_name(leg_type)
-        with open(save_path, "wb") as f:
-            _pkl.dump(model_dict, f)
-
-    @classmethod
-    def from_pickle(
-        cls, save_dir: str, leg_type: _LegType
-    ) -> "SystemPreparationConfig":
-        """
-        Load the configuration from a pickle file.
-
-        Parameters
-        ----------
-        save_dir : str
-            Directory to load the pickle file from.
-
-        leg_type : LegType
-            The type of the leg. Used to decide the name of the pickle
-            file to load.
-
-        Returns
-        -------
-        SystemPreparationConfig
-            Loaded configuration.
-        """
-
-        # Load the dict from the pickle file
-        load_path = save_dir + "/" + cls.get_file_name(leg_type)
-        with open(load_path, "rb") as f:
-            model_dict = _pkl.load(f)
-
-        # Create the model from the dict
-        return cls.parse_obj(model_dict)
-
-    @staticmethod
-    def get_file_name(leg_type: _LegType) -> str:
-        """Get the name of the pickle file for the configuration."""
-        return f"system_preparation_config_{leg_type.name.lower()}.pkl"
+from ..configuration.system_prep_config import (
+    SystemPreparationConfig as _SystemPreparationConfig,
+)
 
 
 def parameterise_input(
@@ -290,7 +46,7 @@ def parameterise_input(
     parameterised_system : _BSS._SireWrappers._system.System
         Parameterised system.
     """
-    cfg = SystemPreparationConfig.from_pickle(input_dir, leg_type)
+    cfg = _SystemPreparationConfig.load(input_dir, leg_type)
 
     print("Parameterising input...")
     # Parameterise the ligand
@@ -382,7 +138,7 @@ def solvate_input(
     solvated_system : _BSS._SireWrappers._system.System
         Solvated system.
     """
-    cfg = SystemPreparationConfig.from_pickle(input_dir, leg_type)
+    cfg = _SystemPreparationConfig.load(input_dir, leg_type)
 
     # Load the parameterised system
     print("Loading parameterised system...")
@@ -476,7 +232,7 @@ def minimise_input(
     minimised_system : _BSS._SireWrappers._system.System
         Minimised system.
     """
-    cfg = SystemPreparationConfig.from_pickle(input_dir, leg_type)
+    cfg = _SystemPreparationConfig.load(input_dir, leg_type)
 
     # Load the solvated system
     print("Loading solvated system...")
@@ -528,7 +284,7 @@ def heat_and_preequil_input(
     preequilibrated_system : _BSS._SireWrappers._system.System
         Pre-Equilibrated system.
     """
-    cfg = SystemPreparationConfig.from_pickle(input_dir, leg_type)
+    cfg = _SystemPreparationConfig.load(input_dir, leg_type)
 
     # Load the minimised system
     print("Loading minimised system...")
@@ -629,7 +385,7 @@ def run_ensemble_equilibration(
     -------
     None
     """
-    cfg = SystemPreparationConfig.from_pickle(input_dir, leg_type)
+    cfg = _SystemPreparationConfig.load(input_dir, leg_type)
 
     # Load the pre-equilibrated system
     print("Loading pre-equilibrated system...")
diff --git a/a3fe/tests/test_system_prep_config.py b/a3fe/tests/test_configuration.py
similarity index 86%
rename from a3fe/tests/test_system_prep_config.py
rename to a3fe/tests/test_configuration.py
index bec1986..70ac16b 100644
--- a/a3fe/tests/test_system_prep_config.py
+++ b/a3fe/tests/test_configuration.py
@@ -1,4 +1,4 @@
-"""Unit and regression tests for the SystemPreparationConfig class."""
+"""Unit and regression tests for Pydantic configuration classes."""
 
 from tempfile import TemporaryDirectory
 
@@ -30,8 +30,8 @@ def test_config_pickle_and_load():
     """Test that the config can be pickled and loaded."""
     with TemporaryDirectory() as dirname:
         config = SystemPreparationConfig()
-        config.save_pickle(dirname, LegType.FREE)
-        config2 = SystemPreparationConfig.from_pickle(dirname, LegType.FREE)
+        config.dump(dirname, LegType.FREE)
+        config2 = SystemPreparationConfig.load(dirname, LegType.FREE)
         assert config == config2
 
 
diff --git a/a3fe/tests/test_run.py b/a3fe/tests/test_run.py
index c2e19e2..b7ccaf4 100644
--- a/a3fe/tests/test_run.py
+++ b/a3fe/tests/test_run.py
@@ -15,7 +15,7 @@
 
 import a3fe as a3
 from a3fe.analyse.detect_equil import dummy_check_equil_multiwindow
-from a3fe.run.system_prep import SystemPreparationConfig
+from a3fe.configuration import SystemPreparationConfig
 
 from . import RUN_SLURM_TESTS, SLURM_PRESENT
 
@@ -195,7 +195,7 @@ def test_parameterisation_free(t4l_calc):
         # We need to save the config to the input directory
         a3.SystemPreparationConfig(
             forcefields={"ligand": "gaff2", "protein": "ff14SB", "water": "tip3p"}
-        ).save_pickle(t4l_calc.input_dir, leg_type)
+        ).dump(t4l_calc.input_dir, leg_type)
         # Parameterise benzene
         free_leg.parameterise_input(slurm=False)
 
@@ -235,7 +235,7 @@ def test_parameterisation_bound(t4l_calc):
         # We need to save the config to the input directory
         a3.SystemPreparationConfig(
             forcefields={"ligand": "gaff2", "protein": "ff14SB", "water": "tip3p"}
-        ).save_pickle(t4l_calc.input_dir, leg_type)
+        ).dump(t4l_calc.input_dir, leg_type)
         # Parameterise benzene
         assert leg_type == a3.LegType.BOUND
         assert bound_leg.leg_type == leg_type
diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml
index aa58b2d..c10ba2b 100644
--- a/devtools/conda-envs/test_env.yaml
+++ b/devtools/conda-envs/test_env.yaml
@@ -23,6 +23,7 @@ dependencies:
   - gromacs
   - pydantic
   - ruff
+  - pyyaml
 
     # Testing
   - pytest
diff --git a/docs/guides.rst b/docs/guides.rst
index 1fc75b5..eab7bdb 100644
--- a/docs/guides.rst
+++ b/docs/guides.rst
@@ -124,7 +124,7 @@ is not killed when you log out).
 Customising Calculations
 *************************
 
-Calculation setup options, including the force fields, lambda schedules, and length of the equilibration steps, can be customised using :class:`a3fe.run.system_prep.SystemPreparationConfig`.
+Calculation setup options, including the force fields, lambda schedules, and length of the equilibration steps, can be customised using :class:`a3fe.configuration.system_preparation.SystemPreparationConfig`.
 For example, to use GAFF2 instead of OFF2 for the small molecule, set this in the config object and pass this to ``calc.setup()``:
 
 .. code-block:: python
diff --git a/environment.yaml b/environment.yaml
index e551c15..e3c4b66 100644
--- a/environment.yaml
+++ b/environment.yaml
@@ -18,3 +18,4 @@ dependencies:
   - ambertools
   - biosimspace>2023.4
   - pydantic
+  - pyyaml
diff --git a/pyproject.toml b/pyproject.toml
index 390d7d8..7d9129d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
      "ipython",
      "pymbar<4",
      "pydantic",
+     "pyyaml",
 ]
 
 # Update the urls once the hosting is set up.