Add the cs_evolution_prior flag to CEST and CPMG experiments (#281)

The `cs_evolution_prior` flag allows for simulations where the CEST or CPMG block occurs after chemical shift evolution. When this flag is set to true in the `[experiment]` section of the `experiment.toml` file, the initial magnetization is populated only for the observed state (`observed_state`), while all other terms are set to zero. This modification ensures accurate boundary conditions by reflecting the magnetization distribution prior to the exchange block.

chemex/configuration/experiment.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import math
+from functools import cached_property
 from typing import Literal, TypeVar
 
 from pydantic import BaseModel, ConfigDict, model_validator
@@ -18,13 +19,25 @@ class ExperimentSettings(BaseModel):
     _key_to_lower = model_validator(mode="before")(key_to_lower)
 
 
-class CpmgSettings(ExperimentSettings):
+class RelaxationSettings(ExperimentSettings):
+    cs_evolution_prior: bool = False
+
+    @cached_property
+    def suffix(self) -> str:
+        return f"_{self.observed_state}" if self.cs_evolution_prior else ""
+
+
+class CpmgSettings(RelaxationSettings):
     even_ncycs: bool = False
 
 
-class CpmgSettingsEvenNcycs(ExperimentSettings):
+class CpmgSettingsEvenNcycs(RelaxationSettings):
     even_ncycs: bool = True
 
 
-class CestSettings(ExperimentSettings):
+class CestSettings(RelaxationSettings):
     sw: float = math.inf
+
+
+class MFCestSettings(RelaxationSettings):
+    sw: float

chemex/experiments/catalog/cest_13c.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -33,9 +34,12 @@ class Cest13CSettings(CestSettings):
     b1_frq: float
     b1_inh_scale: float = 0.1
     b1_inh_res: int = 11
-    observed_state: Literal["a", "b", "c", "d"] = "a"
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[iz_{self.observed_state}]"
 
@@ -88,7 +92,7 @@ def is_reference(metadata: ArrayFloat) -> ArrayBool:
     def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         offsets = data.metadata
 
-        start = spectrometer.get_equilibrium()
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         intensities: dict[float, ArrayFloat] = {}
 

chemex/experiments/catalog/cest_15n.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -34,9 +35,12 @@ class Cest15NSettings(CestSettings):
     b1_frq: float
     b1_inh_scale: float = 0.1
     b1_inh_res: int = 11
-    observed_state: Literal["a", "b", "c", "d"] = "a"
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[iz_{self.observed_state}]"
 
@@ -87,7 +91,7 @@ def is_reference(metadata: ArrayFloat) -> ArrayBool:
     def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         offsets = data.metadata
 
-        start = spectrometer.get_equilibrium()
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         intensities: dict[float, ArrayFloat] = {}
 

chemex/experiments/catalog/cest_15n_cw.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -37,7 +38,11 @@ class Cest15NCwSettings(CestSettings):
     b1_inh_scale: float = 0.1
     b1_inh_res: int = 11
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[iz_{self.observed_state}]"
 
@@ -101,7 +106,9 @@ def is_reference(metadata: ArrayFloat) -> ArrayBool:
     def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         offsets = data.metadata
 
-        start = spectrometer.get_start_magnetization(terms=["iz"], atom="n")
+        start = spectrometer.get_start_magnetization(
+            terms=self.settings.start_terms, atom="n"
+        )
 
         intensities: dict[float, ArrayFloat] = {}
 

chemex/experiments/catalog/cest_15n_tr.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -36,7 +37,19 @@ class Cest15NTrSettings(CestSettings):
     b1_inh_res: int = 11
     antitrosy: bool = False
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        """Start from the TROSY or ANTI-TROSY component.
+
+        TROSY: (2IzSz + Iz) / 2
+        ANTITROSY: (2IzSz - Iz) / 2.
+        """
+        if not self.antitrosy:
+            return [f"2izsz{self.suffix}", f"-iz{self.suffix}"]
+
+        return [f"2izsz{self.suffix}", f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         if self.antitrosy:
             return f"[2izsz_{self.observed_state}] + [iz_{self.observed_state}]"
@@ -95,21 +108,10 @@ class Cest15NTrSequence:
     def is_reference(metadata: ArrayFloat) -> ArrayBool:
         return np.abs(metadata) > OFFSET_REF
 
-    def _get_start(self, spectrometer: Spectrometer) -> ArrayFloat:
-        """Start from the TROSY or ANTI-TROSY component.
-
-        TROSY: (2IzSz - Iz) / 2
-        ANTITROSY: (2IzSz + Iz) / 2.
-        """
-        start = 0.5 * spectrometer.get_start_magnetization(["2izsz", "iz"])
-        if not self.settings.antitrosy:
-            start -= spectrometer.get_start_magnetization(["iz"])
-        return start
-
     def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         offsets = data.metadata
 
-        start = self._get_start(spectrometer)
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         intensities: dict[float, ArrayFloat] = {}
 

chemex/experiments/catalog/cest_1hn_ap.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -33,13 +34,13 @@ class Cest1HnApSettings(CestSettings):
     b1_frq: float
     b1_inh_scale: float = 0.1
     b1_inh_res: int = 11
-    observed_state: Literal["a", "b", "c", "d"] = "a"
+    cs_evolution_prior: bool = True
 
-    @property
-    def start(self) -> list[str]:
-        return [f"2izsz_{self.observed_state}"]
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"2izsz{self.suffix}"]
 
-    @property
+    @cached_property
     def detection(self) -> str:
         return f"[2izsz_{self.observed_state}]"
 
@@ -90,7 +91,7 @@ def is_reference(metadata: ArrayFloat) -> ArrayBool:
     def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         offsets = data.metadata
 
-        start = spectrometer.get_start_magnetization(terms=self.settings.start)
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         intensities: dict[float, ArrayFloat] = {}
 

chemex/experiments/catalog/coscest_13c.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -9,7 +10,7 @@
 from chemex.configuration.base import ExperimentConfiguration, ToBeFitted
 from chemex.configuration.conditions import ConditionsWithValidations
 from chemex.configuration.data import CestDataSettings
-from chemex.configuration.experiment import CestSettings
+from chemex.configuration.experiment import MFCestSettings
 from chemex.containers.data import Data
 from chemex.containers.dataset import load_relaxation_dataset
 from chemex.experiments.factories import Creators, factories
@@ -28,19 +29,22 @@
 OFFSET_REF = 1e4
 
 
-class CosCest13CSettings(CestSettings):
+class CosCest13CSettings(MFCestSettings):
     name: Literal["coscest_13c"]
     time_t1: float
     time_equil: float = 0.0
     carrier: float
-    sw: float
     cos_n: int
     cos_res: int = 10
     b1_frq: float
     b1_inh_scale: float = 0.1
     b1_inh_res: int = 11
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[iz_{self.observed_state}]"
 

chemex/experiments/catalog/coscest_1hn_ip_ap.py

@@ -9,7 +9,7 @@
 from chemex.configuration.base import ExperimentConfiguration, ToBeFitted
 from chemex.configuration.conditions import ConditionsWithValidations
 from chemex.configuration.data import CestDataSettingsNoRef
-from chemex.configuration.experiment import CestSettings
+from chemex.configuration.experiment import MFCestSettings
 from chemex.containers.data import Data
 from chemex.containers.dataset import load_relaxation_dataset
 from chemex.experiments.factories import Creators, factories
@@ -27,11 +27,10 @@
 OFFSET_REF = 1e4
 
 
-class CosCest1HnIpApSettings(CestSettings):
+class CosCest1HnIpApSettings(MFCestSettings):
     name: Literal["coscest_1hn_ip_ap"]
     time_t1: float
     carrier: float
-    sw: float
     cos_n: int
     cos_res: int = 10
     d1: float

chemex/experiments/catalog/cpmg_13c_ip.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -32,11 +33,15 @@ class Cpmg13CIpSettings(CpmgSettings):
     pw90: float
     time_equil: float = 0.0
 
-    @property
+    @cached_property
     def t_neg(self) -> float:
         return -2.0 * self.pw90 / np.pi
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[iz_{self.observed_state}]"
 
@@ -104,7 +109,7 @@ def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         p180pmx = 0.5 * (p180[0] + p180[2])  # +/- phase cycling
 
         # Getting the starting magnetization
-        start = spectrometer.get_equilibrium()
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         # Calculating the instensities as a function of ncyc
         part1 = d_neg @ p90[0] @ start

chemex/experiments/catalog/cpmg_13co_ap.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -34,15 +35,19 @@ class Cpmg13CoApSettings(CpmgSettings):
     refocusing: bool = False
     taucc: float = 9.09e-3
 
-    @property
+    @cached_property
     def t_neg(self) -> float:
         return -2.0 * self.pw90 / np.pi
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"2izsz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[2izsz_{self.observed_state}]"
 
-    @property
+    @cached_property
     def even_ncycs(self) -> bool:
         return self.refocusing
 
@@ -114,7 +119,7 @@ def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         perfect180x = spectrometer.perfect180_i[0]
 
         # Getting the starting magnetization
-        start = spectrometer.get_start_magnetization(["2izsz"])
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         # Calculate the flip block
         if self.settings.refocusing:

chemex/experiments/catalog/cpmg_15n_ip.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Literal
 
 import numpy as np
@@ -32,11 +33,15 @@ class Cpmg15NIpSettings(CpmgSettings):
     pw90: float
     time_equil: float = 0.0
 
-    @property
+    @cached_property
     def t_neg(self) -> float:
         return -2.0 * self.pw90 / np.pi
 
-    @property
+    @cached_property
+    def start_terms(self) -> list[str]:
+        return [f"iz{self.suffix}"]
+
+    @cached_property
     def detection(self) -> str:
         return f"[iz_{self.observed_state}]"
 
@@ -104,7 +109,7 @@ def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
         p180pmx = 0.5 * (p180[0] + p180[2])  # +/- phase cycling
 
         # Getting the starting magnetization
-        start = spectrometer.get_equilibrium()
+        start = spectrometer.get_start_magnetization(self.settings.start_terms)
 
         # Calculating the instensities as a function of ncyc
         part1 = d_neg @ p90[0] @ start