a b/aiagents4pharma/talk2biomodels/tools/load_arguments.py
1
#!/usr/bin/env python3
2
3
"""
4
A utility module for defining the dataclasses
5
for the arguments to set up initial settings
6
before the experiment is run.
7
"""
8
9
import logging
10
from dataclasses import dataclass
11
from typing import Union, List, Optional, Annotated
12
from pydantic import Field
13
import basico
14
15
# Initialize logger
16
logging.basicConfig(level=logging.INFO)
17
logger = logging.getLogger(__name__)
18
19
@dataclass
20
class TimeData:
21
    """
22
    Dataclass for storing the time data.
23
    """
24
    duration: Union[int, float] = Field(
25
        description="Duration of the simulation",
26
        default=100)
27
    interval: Union[int, float] = Field(
28
        description="The interval is the time step or"
29
        " the step size of the simulation. It is unrelated"
30
        " to the step size of species concentration and parameter values.",
31
        default=100)
32
33
@dataclass
34
class SpeciesInitialData:
35
    """
36
    Dataclass for storing the species initial data.
37
    """
38
    species_name: List[str] = Field(
39
        description="List of species whose initial concentration is to be set."
40
        " This does not include species that reoccur or the species whose"
41
        " concentration is to be determined/observed at the end of the experiment."
42
        " Do not hallucinate the species name.",
43
        default=[])
44
    species_concentration: List[Union[int, float]] = Field(
45
        description="List of initial concentrations of species."
46
        " This does not include species that reoccur or the species whose"
47
        " concentration is to be determined/observed at the end of the experiment."
48
        " Do not hallucinate the species concentration.",
49
        default=[])
50
51
@dataclass
52
class TimeSpeciesNameConcentration:
53
    """
54
    Dataclass for storing the time,
55
    species name, and concentration data.
56
    """
57
    time: Union[int, float] = Field(description="time point where the event occurs")
58
    species_name: str = Field(description="species name")
59
    species_concentration: Union[int, float] = Field(
60
        description="species concentration at the time point")
61
62
@dataclass
63
class ReocurringData:
64
    """
65
    Dataclass for species that reoccur. In other words,
66
    the concentration of the species resets to a certain
67
    value after a certain time interval.
68
    """
69
    data: List[TimeSpeciesNameConcentration] = Field(
70
        description="List of time, name, and concentration data"
71
                    " of species or parameters that reoccur",
72
                    default=[])
73
74
@dataclass
75
class ArgumentData:
76
    """
77
    Dataclass for storing the argument data.
78
    """
79
    experiment_name: Annotated[str, "An AI assigned _ separated name of"
80
                                    " the experiment based on human query"
81
                                    " and the context of the experiment."
82
                                    " This must be set before the experiment is run."]
83
    time_data: Optional[TimeData] = Field(
84
        description="time data",
85
        default=None)
86
    species_to_be_analyzed_before_experiment: Optional[SpeciesInitialData] = Field(
87
        description="Data of species whose initial concentration"
88
        " is to be set before the experiment. This does not include"
89
        " species that reoccur or the species whose concentration"
90
        " is to be determined at the end of the experiment.",
91
        default=None)
92
    reocurring_data: Optional[ReocurringData] = Field(
93
        description="List of concentration and time data of species that"
94
        " reoccur. For example, a species whose concentration resets"
95
        " to a certain value after a certain time interval.",
96
        default=None)
97
98
def add_rec_events(model_object, reocurring_data):
99
    """
100
    Add reocurring events to the model.
101
102
    Args:
103
        model_object: The model object.
104
        reocurring_data: The reocurring data.
105
106
    Returns:
107
        None
108
    """
109
    for row in reocurring_data.data:
110
        tp, sn, sc = row.time, row.species_name, row.species_concentration
111
        basico.add_event(f'{sn}_{tp}',
112
                            f'Time > {tp}',
113
                            [[sn, str(sc)]],
114
                            model=model_object.copasi_model)