abstochkin.agentstatedata

  1"""
  2Class for storing the state of all agents of a certain species during an
  3AbStochKin simulation.
  4"""
  5
  6#  Copyright (c) 2024-2025, Alex Plakantonakis.
  7#
  8#  This program is free software: you can redistribute it and/or modify
  9#  it under the terms of the GNU General Public License as published by
 10#  the Free Software Foundation, either version 3 of the License, or
 11#  (at your option) any later version.
 12#
 13#  This program is distributed in the hope that it will be useful,
 14#  but WITHOUT ANY WARRANTY; without even the implied warranty of
 15#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 16#  GNU General Public License for more details.
 17#
 18#  You should have received a copy of the GNU General Public License
 19#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 20
 21from copy import deepcopy
 22from dataclasses import dataclass, field
 23
 24import numpy as np
 25
 26
 27@dataclass
 28class AgentStateData:
 29    """
 30    Class for storing the state of all agents of a certain species during an
 31    AbStochKin simulation.
 32
 33    Attributes
 34    ----------
 35    p_init : int
 36        The initial population size of the species whose data is
 37        represented in an `AgentStateData` object.
 38    max_agents : int
 39        The maximum number of agents for the species whose data
 40        is represented in an `AgentStateData` object.
 41    reps : int
 42        The number of times the AbStochKin algorithm will repeat a simulation.
 43        This will be the length of the `asv` list.
 44    fill_state: int
 45
 46    asv_ini : numpy.ndarray
 47        Agent-State Vector (asv) is a species-specific 2-row vector to monitor
 48        agent state according to Markov's property. This array is the initial
 49        asv, i.e., at `t=0`. The array shape is `(2, max_agents)`
 50    asv : list of numpy.ndarray
 51        A list of length `reps` with copies of `asv_ini`. Each simulation run
 52        uses its corresponding entry in `asv` to monitor the state of all
 53        agents.
 54    """
 55
 56    p_init: int  # initial population size
 57    max_agents: int  # maximum number of agents represented in asv
 58    reps: int  # number of times simulation is repeated
 59    fill_state: int
 60
 61    asv_ini: np.ndarray = field(init=False, default_factory=lambda: np.array([]))
 62    asv: list[np.ndarray] = field(init=False, default_factory=lambda: list(np.array([])))
 63
 64    def __post_init__(self):
 65        # Set up initial (t=0) agent-state vector (asv):
 66        self.asv_ini = np.concatenate(
 67            (np.ones(shape=(2, self.p_init), dtype=np.int8),
 68             np.full(shape=(2, self.max_agents - self.p_init),
 69                     fill_value=self.fill_state,
 70                     dtype=np.int8)),
 71            axis=1
 72        )
 73
 74        # Set up separate copy of the initial `asv` for each repetition of the
 75        # algorithm to facilitate parallelization of ensemble runs.
 76        self.asv = [deepcopy(self.asv_ini) for _ in range(self.reps)]
 77
 78    def apply_markov_property(self, r: int):
 79        """
 80        The future state of the system depends only on its current state.
 81        This method is called at the end of each time step in an AbStochKin
 82        simulation. Therefore, the new agent-state vector becomes the
 83        current state.
 84        """
 85        self.asv[r][0, :] = self.asv[r][1, :]
 86
 87    def cleanup_asv(self):
 88        """ Empty the contents of the array `asv`. """
 89        self.asv = list(np.array([]))
 90
 91    def get_vals_o1(self,
 92                    r: int,
 93                    stream: np.random.Generator,
 94                    p_vals: np.ndarray,
 95                    state: int = 1):
 96        """
 97        Get random values in [0,1) at a given time step for agents of a given
 98        state. Agents of other states have a value of zero.
 99
100        Get probability values at a given time step for agents of a given state.
101        Agents of other states have a transition probability of zero.
102
103        Notes
104        -----
105        Note that only elements of the `asv` that have the same state in the
106        previous and current time steps are considered. This is to ensure that
107        agents that have already transitioned to a different state in the
108        current time step are not reconsidered for a possible transition.
109        """
110        nonzero_elems = np.all(self.asv[r] == state, axis=0)
111        final_rand_nums = stream.random(self.max_agents) * nonzero_elems
112        final_p_vals = p_vals * nonzero_elems
113
114        return final_rand_nums, final_p_vals
115
116    def get_vals_o2(self,
117                    other,
118                    r: int,
119                    stream: np.random.Generator,
120                    p_vals: np.ndarray,
121                    state: int = 1):
122        """
123        Get random values in [0,1) at a given time step for interactions between
124        agents of a given state. Agents of other states have a value of zero.
125
126        Get probability values at a given time step for interactions between
127        agents of a given state. Interactions of agents in other states
128        have a transition probability of zero.
129
130        Notes
131        -----
132        Note that only elements of the `asv` that have the same state in the
133        previous and current time steps are considered. This is to ensure that
134        agents that have already transitioned to a different state in the
135        current time step are not reconsidered for a possible transition.
136        """
137        nonzero_rows = np.all(self.asv[r] == state, axis=0).reshape(-1, 1)
138        nonzero_cols = np.all(other.asv[r] == state, axis=0).reshape(1, -1)
139
140        rand_nums = stream.random(size=(self.max_agents, other.max_agents))
141
142        final_rand_nums = rand_nums * nonzero_rows * nonzero_cols
143        final_p_vals = p_vals * nonzero_rows * nonzero_cols
144
145        return final_rand_nums, final_p_vals
146
147    def __str__(self):
148        return f"Agent-State Vector with \n" \
149               f"Initial population size: {self.p_init}\n" \
150               f"Maximum number of agents: {self.max_agents}\n" \
151               f"Repeat simulation {self.reps} times\n" \
152               f"Fill state: {self.fill_state}"