Skip to content

enzax.kinetic_model

Module containing enzax's definition of a kinetic model.

KineticModel

Structural information about a kinetic model.

Source code in enzax/kinetic_model.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class KineticModel(eqx.Module):
    """Structural information about a kinetic model."""

    stoichiometry: dict[str, dict[str, float]] = eqx.field(static=True)
    species: list[str] = eqx.field(static=True)
    reactions: list[str] = eqx.field(static=True)
    balanced_species: list[str] = eqx.field(static=True)
    unbalanced_species: list[str] = eqx.field(static=True, init=False)
    species_to_dgf_ix: NDArray[np.int16] = eqx.field(
        static=True, default=slice(None)
    )
    balanced_species_ix: NDArray[np.int16] = eqx.field(static=True, init=False)
    unbalanced_species_ix: NDArray[np.int16] = eqx.field(
        static=True, init=False
    )
    S: NDArray[np.float64] = eqx.field(static=True, init=False)

    def __post_init__(self, species_to_dgf_ix=None):
        self.unbalanced_species = [
            s for s in self.species if s not in self.balanced_species
        ]
        self.balanced_species_ix = np.array(
            [get_ix_from_list(s, self.species) for s in self.balanced_species],
            dtype=np.int16,
        )
        self.unbalanced_species_ix = np.array(
            [
                get_ix_from_list(s, self.species)
                for s in self.unbalanced_species
            ],
            dtype=np.int16,
        )
        S = np.zeros(shape=(len(self.species), len(self.reactions)))
        for ix_reaction, reaction in enumerate(self.reactions):
            for species_i, coeff in self.stoichiometry[reaction].items():
                ix_species = get_ix_from_list(species_i, self.species)
                S[ix_species, ix_reaction] = coeff
        self.S = S.astype(np.float64)

    def tree_flatten(self):
        children = (
            self.stoichiometry,
            self.species,
            self.reactions,
            self.balanced_species,
            self.species_to_dgf_ix,
        )
        aux_data = None
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children)

    def get_conc(self, balanced, log_unbalanced):
        conc = jnp.zeros(self.S.shape[0])
        conc = conc.at[self.balanced_species_ix].set(balanced)
        conc = conc.at[self.unbalanced_species_ix].set(jnp.exp(log_unbalanced))
        return conc

    @abstractmethod
    def flux(
        self,
        conc_balanced: Float[Array, " n_balanced"],
        parameters: PyTree,
    ) -> Float[Array, " n"]: ...

    def dcdt(
        self,
        conc: Float[Array, " n_balanced"],
        parameters: PyTree,
    ) -> Float[Array, " n_balanced"]:
        """Get the rate of change of balanced species concentrations.

        :param conc: a one dimensional array of positive floats representing concentrations of balanced species. Must have same size as self.structure.ix_balanced

        :param parameters: A PyTree of parameters.

        :return: a one dimensional array of floats representing the rate of change of balanced species concentrations. Has same size as self.structure.ix_balanced.
        """  # Noqa: E501
        v = self.flux(conc, parameters)
        sv = self.S @ v
        return jnp.array(sv[self.balanced_species_ix])

    def __call__(self, t, y, parameters):
        return self.dcdt(y, parameters)

dcdt(conc, parameters)

Get the rate of change of balanced species concentrations.

Parameters:

Name Type Description Default
conc Float[Array, ' n_balanced']

a one dimensional array of positive floats representing concentrations of balanced species. Must have same size as self.structure.ix_balanced

 required
parameters PyTree

A PyTree of parameters.

 required

Returns:

Type Description
Float[Array, ' n_balanced']

a one dimensional array of floats representing the rate of change of balanced species concentrations. Has same size as self.structure.ix_balanced.
Source code in enzax/kinetic_model.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def dcdt(
    self,
    conc: Float[Array, " n_balanced"],
    parameters: PyTree,
) -> Float[Array, " n_balanced"]:
    """Get the rate of change of balanced species concentrations.

    :param conc: a one dimensional array of positive floats representing concentrations of balanced species. Must have same size as self.structure.ix_balanced

    :param parameters: A PyTree of parameters.

    :return: a one dimensional array of floats representing the rate of change of balanced species concentrations. Has same size as self.structure.ix_balanced.
    """  # Noqa: E501
    v = self.flux(conc, parameters)
    sv = self.S @ v
    return jnp.array(sv[self.balanced_species_ix])

KineticModelSbml

Source code in enzax/kinetic_model.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
class KineticModelSbml(KineticModel):
    sym_module: Any = eqx.field(static=True, default=None)

    def flux(
        self,
        conc_balanced: Float[Array, " n_balanced"],
        parameters,
    ) -> Float[Array, " n"]:
        flux = jnp.array(
            self.sym_module(
                **parameters,
                **dict(zip(self.balanced_species, conc_balanced)),
            )
        )
        return flux

RateEquationModel

A kinetic model that specifies its fluxes using RateEquation objects.

Source code in enzax/kinetic_model.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
class RateEquationModel(KineticModel):
    """A kinetic model that specifies its fluxes using RateEquation objects."""

    rate_equations: list[RateEquation] = eqx.field(
        static=True, default_factory=list
    )

    def flux(
        self,
        conc_balanced: Float[Array, " n_balanced"],
        parameters: PyTree,
    ) -> Float[Array, " n"]:
        """Get fluxes from balanced species concentrations.

        :param conc_balanced: a one dimensional array of positive floats representing concentrations of balanced species. Must have same size as self.structure.ix_balanced

        :return: a one dimensional array of (possibly negative) floats representing reaction fluxes. Has same size as number of columns of self.structure.S.

        """  # Noqa: E501
        conc = self.get_conc(conc_balanced, parameters["log_conc_unbalanced"])
        flux_list = []
        for reaction_ix, (reaction_id, rate_equation) in enumerate(
            zip(self.reactions, self.rate_equations)
        ):
            ipt = rate_equation.get_input(
                parameters=parameters,
                reaction_id=reaction_id,
                reaction_stoichiometry=self.S[:, reaction_ix],
                species_to_dgf_ix=self.species_to_dgf_ix,
            )
            flux_list.append(rate_equation(conc, ipt))
        return jnp.array(flux_list)

flux(conc_balanced, parameters)

Get fluxes from balanced species concentrations.

Parameters:

Name Type Description Default
conc_balanced Float[Array, ' n_balanced']

a one dimensional array of positive floats representing concentrations of balanced species. Must have same size as self.structure.ix_balanced

 required

Returns:

Type Description
Float[Array, ' n']

a one dimensional array of (possibly negative) floats representing reaction fluxes. Has same size as number of columns of self.structure.S.
Source code in enzax/kinetic_model.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def flux(
    self,
    conc_balanced: Float[Array, " n_balanced"],
    parameters: PyTree,
) -> Float[Array, " n"]:
    """Get fluxes from balanced species concentrations.

    :param conc_balanced: a one dimensional array of positive floats representing concentrations of balanced species. Must have same size as self.structure.ix_balanced

    :return: a one dimensional array of (possibly negative) floats representing reaction fluxes. Has same size as number of columns of self.structure.S.

    """  # Noqa: E501
    conc = self.get_conc(conc_balanced, parameters["log_conc_unbalanced"])
    flux_list = []
    for reaction_ix, (reaction_id, rate_equation) in enumerate(
        zip(self.reactions, self.rate_equations)
    ):
        ipt = rate_equation.get_input(
            parameters=parameters,
            reaction_id=reaction_id,
            reaction_stoichiometry=self.S[:, reaction_ix],
            species_to_dgf_ix=self.species_to_dgf_ix,
        )
        flux_list.append(rate_equation(conc, ipt))
    return jnp.array(flux_list)