"""The glotaran generator module."""
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import Any
from typing import TypedDict
from typing import cast
from glotaran.builtin.io.yml.utils import write_dict
from glotaran.builtin.megacomplexes.decay import DecayParallelMegacomplex
from glotaran.builtin.megacomplexes.decay import DecaySequentialMegacomplex
from glotaran.builtin.megacomplexes.spectral import SpectralMegacomplex
from glotaran.model import Model
if TYPE_CHECKING:
from collections.abc import Callable
[docs]
def _generate_decay_model(
*, nr_compartments: int, irf: bool, spectral: bool, decay_type: str
) -> dict[str, Any]:
"""Generate a decay model dictionary.
Parameters
----------
nr_compartments : int
The number of compartments.
irf : bool
Whether to add a gaussian irf.
spectral : bool
Whether to add a spectral model.
decay_type : str
The type of the decay
Returns
-------
dict[str, Any]
The generated model dictionary.
"""
compartments = [f"species_{i + 1}" for i in range(nr_compartments)]
rates = [f"rates.species_{i + 1}" for i in range(nr_compartments)]
model: dict[str, Any] = {
"megacomplex": {
f"megacomplex_{decay_type}_decay": {
"type": f"decay-{decay_type}",
"compartments": compartments,
"rates": rates,
},
},
"dataset": {"dataset_1": {"megacomplex": [f"megacomplex_{decay_type}_decay"]}},
}
if spectral:
model["megacomplex"] |= {
"megacomplex_spectral": {
"type": "spectral",
"shape": {
compartment: f"shape_species_{i + 1}"
for i, compartment in enumerate(compartments)
},
}
}
model["shape"] = {
f"shape_species_{i + 1}": {
"type": "skewed-gaussian",
"amplitude": f"shapes.species_{i + 1}.amplitude",
"location": f"shapes.species_{i + 1}.location",
"width": f"shapes.species_{i + 1}.width",
"skewness": f"shapes.species_{i + 1}.skewness",
}
for i in range(nr_compartments)
}
model["dataset"]["dataset_1"] |= {
"global_megacomplex": ["megacomplex_spectral"],
"spectral_axis_inverted": True,
"spectral_axis_scale": 1e7,
}
if irf:
model["dataset"]["dataset_1"] |= {"irf": "gaussian_irf"}
model["irf"] = {
"gaussian_irf": {"type": "gaussian", "center": "irf.center", "width": "irf.width"},
}
return model
[docs]
def generate_parallel_decay_model(
*, nr_compartments: int = 1, irf: bool = False
) -> dict[str, Any]:
"""Generate a parallel decay model dictionary.
Parameters
----------
nr_compartments : int
The number of compartments.
irf : bool
Whether to add a gaussian irf.
Returns
-------
dict[str, Any]
The generated model dictionary.
"""
return _generate_decay_model(
nr_compartments=nr_compartments, irf=irf, spectral=False, decay_type="parallel"
)
[docs]
def generate_parallel_spectral_decay_model(
*, nr_compartments: int = 1, irf: bool = False
) -> dict[str, Any]:
"""Generate a parallel spectral decay model dictionary.
Parameters
----------
nr_compartments : int
The number of compartments.
irf : bool
Whether to add a gaussian irf.
Returns
-------
dict[str, Any]
The generated model dictionary.
"""
return _generate_decay_model(
nr_compartments=nr_compartments, irf=irf, spectral=True, decay_type="parallel"
)
[docs]
def generate_sequential_decay_model(
*, nr_compartments: int = 1, irf: bool = False
) -> dict[str, Any]:
"""Generate a sequential decay model dictionary.
Parameters
----------
nr_compartments : int
The number of compartments.
irf : bool
Whether to add a gaussian irf.
Returns
-------
dict[str, Any]
The generated model dictionary.
"""
return _generate_decay_model(
nr_compartments=nr_compartments, irf=irf, spectral=False, decay_type="sequential"
)
[docs]
def generate_sequential_spectral_decay_model(
*, nr_compartments: int = 1, irf: bool = False
) -> dict[str, Any]:
"""Generate a sequential spectral decay model dictionary.
Parameters
----------
nr_compartments : int
The number of compartments.
irf : bool
Whether to add a gaussian irf.
Returns
-------
dict[str, Any]
The generated model dictionary.
"""
return _generate_decay_model(
nr_compartments=nr_compartments, irf=irf, spectral=True, decay_type="sequential"
)
generators: dict[str, Callable] = {
"decay_parallel": generate_parallel_decay_model,
"spectral_decay_parallel": generate_parallel_spectral_decay_model,
"decay_sequential": generate_sequential_decay_model,
"spectral_decay_sequential": generate_sequential_spectral_decay_model,
}
available_generators: list[str] = list(generators.keys())
[docs]
class GeneratorArguments(TypedDict, total=False):
"""Arguments used by ``generate_model`` and ``generate_model``.
Parameters
----------
nr_compartments : int
The number of compartments.
irf : bool
Whether to add a gaussian irf.
See Also
--------
generate_model
generate_model_yml
"""
nr_compartments: int
irf: bool
[docs]
def generate_model(*, generator_name: str, generator_arguments: GeneratorArguments) -> Model:
"""Generate a model.
Parameters
----------
generator_name : str
The generator to use.
generator_arguments : GeneratorArguments
Arguments for the generator.
Returns
-------
Model
The generated model
See Also
--------
generate_parallel_decay_model
generate_parallel_spectral_decay_model
generate_sequential_decay_model
generate_sequential_spectral_decay_model
Raises
------
ValueError
Raised when an unknown generator is specified.
"""
if generator_name not in generators:
msg = (
f"Unknown model generator '{generator_name}'. "
f"Known generators are: {list(generators.keys())}"
)
raise ValueError(msg)
model = generators[generator_name](**generator_arguments)
return Model.create_class_from_megacomplexes(
[DecayParallelMegacomplex, DecaySequentialMegacomplex, SpectralMegacomplex]
)(**model)
[docs]
def generate_model_yml(*, generator_name: str, generator_arguments: GeneratorArguments) -> str:
"""Generate a model as yml string.
Parameters
----------
generator_name : str
The generator to use.
generator_arguments : GeneratorArguments
Arguments for the generator.
Returns
-------
str
The generated model yml string.
See Also
--------
generate_parallel_decay_model
generate_parallel_spectral_decay_model
generate_sequential_decay_model
generate_sequential_spectral_decay_model
Raises
------
ValueError
Raised when an unknown generator is specified.
"""
if generator_name not in generators:
msg = (
f"Unknown model generator '{generator_name}'. "
f"Known generators are: {list(generators.keys())}"
)
raise ValueError(msg)
model = generators[generator_name](**generator_arguments)
return cast("str", write_dict(model))