Skip to content

Commit

Permalink
Merge branch 'main' into hjafari/MIC-5536_drop_python_3.9_changes
Browse files Browse the repository at this point in the history
  • Loading branch information
hussain-jafari committed Nov 13, 2024
2 parents 02d12be + 881d6b8 commit 257fb3d
Show file tree
Hide file tree
Showing 30 changed files with 116 additions and 51 deletions.
8 changes: 7 additions & 1 deletion CHANGELOG.rst
Original file line number Diff line number Diff line change
@@ -1,3 +1,9 @@
**3.2.1 - 11/13/24**

- Fix mypy errors in vivarium/framework/results/context.py
- Modernize type hinting
- Remove unnecessary "from future import annotation" imports

**3.2.0 - 11/12/24**

- Feature: Supports passing callables directly when building lookup tables
Expand All @@ -9,7 +15,7 @@
- Testing: Adds coverage for example DiseaseModel
- Refactor: Converts resource module into a package
- Refactor: Converts values module into a package
= Refactor: Simplifies code to allow Managers to create columns
- Refactor: Simplifies code to allow Managers to create columns
- Refactor: Converts ResourceManager __iter__ to a well-named instance method
- Refactor: Creates ResourceTypes for each type of resource
- Refactor: Makes Pipeline and RandomnessStream inherit from Resource
Expand Down
26 changes: 13 additions & 13 deletions docs/source/concepts/results.rst
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,7 @@ A couple other more specific and commonly used observations are provided as well
that gathers new results and concatenates them to any existing results.

Ideally, all concrete classes should inherit from the
:class:`BaseObservation <vivarium.framework.results.observation.BaseObservation>`
:class:`Observation <vivarium.framework.results.observation.Observation>`
abstract base class, which contains the common attributes between observation types:

.. list-table:: **Common Observation Attributes**
Expand All @@ -312,40 +312,40 @@ abstract base class, which contains the common attributes between observation ty

* - Attribute
- Description
* - | :attr:`name <vivarium.framework.results.observation.BaseObservation.name>`
* - | :attr:`name <vivarium.framework.results.observation.Observation.name>`
- | Name of the observation. It will also be the name of the output results file
| for this particular observation.
* - | :attr:`pop_filter <vivarium.framework.results.observation.BaseObservation.pop_filter>`
* - | :attr:`pop_filter <vivarium.framework.results.observation.Observation.pop_filter>`
- | A Pandas query filter string to filter the population down to the simulants
| who should be considered for the observation.
* - | :attr:`when <vivarium.framework.results.observation.BaseObservation.when>`
* - | :attr:`when <vivarium.framework.results.observation.Observation.when>`
- | Name of the lifecycle phase the observation should happen. Valid values are:
| "time_step__prepare", "time_step", "time_step__cleanup", or "collect_metrics".
* - | :attr:`results_initializer <vivarium.framework.results.observation.BaseObservation.results_initializer>`
* - | :attr:`results_initializer <vivarium.framework.results.observation.Observation.results_initializer>`
- | Method or function that initializes the raw observation results
| prior to starting the simulation. This could return, for example, an empty
| DataFrame or one with a complete set of stratifications as the index and
| all values set to 0.0.
* - | :attr:`results_gatherer <vivarium.framework.results.observation.BaseObservation.results_gatherer>`
* - | :attr:`results_gatherer <vivarium.framework.results.observation.Observation.results_gatherer>`
- | Method or function that gathers the new observation results.
* - | :attr:`results_updater <vivarium.framework.results.observation.BaseObservation.results_updater>`
* - | :attr:`results_updater <vivarium.framework.results.observation.Observation.results_updater>`
- | Method or function that updates existing raw observation results with newly
| gathered results.
* - | :attr:`results_formatter <vivarium.framework.results.observation.BaseObservation.results_formatter>`
* - | :attr:`results_formatter <vivarium.framework.results.observation.Observation.results_formatter>`
- | Method or function that formats the raw observation results.
* - | :attr:`stratifications <vivarium.framework.results.observation.BaseObservation.stratifications>`
* - | :attr:`stratifications <vivarium.framework.results.observation.Observation.stratifications>`
- | Optional tuple of column names for the observation to stratify by.
* - | :attr:`to_observe <vivarium.framework.results.observation.BaseObservation.to_observe>`
* - | :attr:`to_observe <vivarium.framework.results.observation.Observation.to_observe>`
- | Method or function that determines whether to perform an observation on this Event.

The **BaseObservation** also contains the
:meth:`observe <vivarium.framework.results.observation.BaseObservation.observe>`
The **Observation** also contains the
:meth:`observe <vivarium.framework.results.observation.Observation.observe>`
method which is called at each :ref:`event <event_concept>` and :ref:`time step <time_concept>`
to determine whether or not the observation should be recorded, and if so, gathers
the results and stores them in the results system.

.. note::
All four observation types discussed above inherit from the **BaseObservation**
All four observation types discussed above inherit from the **Observation**
abstract base class. What differentiates them are the assigned attributes
(e.g. defining the **results_updater** to be an adding method for the
**AddingObservation**) or adding other attributes as necessary (e.g.
Expand Down
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,6 @@ exclude = [
'src/vivarium/framework/lookup/manager.py',
'src/vivarium/framework/population/manager.py',
'src/vivarium/framework/population/population_view.py',
'src/vivarium/framework/results/context.py',
'src/vivarium/framework/results/interface.py',
'src/vivarium/framework/results/manager.py',
'src/vivarium/framework/results/observer.py',
Expand Down
7 changes: 4 additions & 3 deletions src/vivarium/component.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
simulations.
"""
from __future__ import annotations

import re
import warnings
from abc import ABC
Expand All @@ -30,8 +32,7 @@
from vivarium.framework.event import Event
from vivarium.framework.lookup import LookupTable
from vivarium.framework.population import PopulationView, SimulantData
from vivarium.framework.randomness import RandomnessStream
from vivarium.framework.values import Pipeline
from vivarium.framework.resource import Resource
from vivarium.types import LookupTableData

DEFAULT_EVENT_PRIORITY = 5
Expand Down Expand Up @@ -237,7 +238,7 @@ def columns_required(self) -> list[str] | None:
@property
def initialization_requirements(
self,
) -> list[str | Pipeline | RandomnessStream]:
) -> list[str | Resource]:
"""A list containing the columns, pipelines, and randomness streams
required by this component's simulant initializer."""
return []
Expand Down
7 changes: 4 additions & 3 deletions src/vivarium/examples/disease_model/risk.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
# mypy: ignore-errors
from __future__ import annotations

from typing import TYPE_CHECKING, Any

import pandas as pd
Expand All @@ -7,8 +9,7 @@

if TYPE_CHECKING:
from vivarium.framework.engine import Builder
from vivarium.framework.randomness import RandomnessStream
from vivarium.framework.values import Pipeline
from vivarium.framework.resource import Resource


class Risk(Component):
Expand All @@ -31,7 +32,7 @@ def columns_created(self) -> list[str]:
return [self.propensity_column]

@property
def initialization_requirements(self) -> list[str | Pipeline | RandomnessStream]:
def initialization_requirements(self) -> list[str | Resource]:
return [self.randomness]

#####################
Expand Down
1 change: 1 addition & 0 deletions src/vivarium/framework/event.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
:ref:`concept note <event_concept>`.
"""
from __future__ import annotations

from collections.abc import Callable
from datetime import datetime, timedelta
Expand Down
1 change: 1 addition & 0 deletions src/vivarium/framework/lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
The tools here also allow for introspection of the simulation life cycle.
"""
from __future__ import annotations

import functools
import textwrap
Expand Down
1 change: 1 addition & 0 deletions src/vivarium/framework/logging/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
=====================
"""
from __future__ import annotations

import loguru
from loguru import logger
Expand Down
2 changes: 2 additions & 0 deletions src/vivarium/framework/logging/utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
This module contains utilities for configuring logging.
"""
from __future__ import annotations

import logging
import sys
from pathlib import Path
Expand Down
2 changes: 2 additions & 0 deletions src/vivarium/framework/lookup/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable, Sequence
from datetime import datetime
Expand Down
2 changes: 2 additions & 0 deletions src/vivarium/framework/population/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
:ref:`population management system <population_concept>`.
"""
from __future__ import annotations

from collections.abc import Callable, Iterable, Sequence
from dataclasses import dataclass
from types import MethodType
Expand Down
2 changes: 2 additions & 0 deletions src/vivarium/framework/randomness/index_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
"""

from __future__ import annotations

from datetime import datetime
from typing import Any

Expand Down
2 changes: 2 additions & 0 deletions src/vivarium/framework/randomness/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
=========================
"""
from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING

Expand Down
2 changes: 2 additions & 0 deletions src/vivarium/framework/randomness/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@
"""

from __future__ import annotations

import hashlib
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, TypeVar
Expand Down
34 changes: 23 additions & 11 deletions src/vivarium/framework/results/context.py
Original file line number Diff line number Diff line change
@@ -1,28 +1,29 @@
# mypy: ignore-errors
"""
===============
Results Context
===============
"""

from __future__ import annotations

from collections import defaultdict
from collections.abc import Callable, Generator
from typing import Type
from typing import Any

import pandas as pd
from pandas.core.groupby.generic import DataFrameGroupBy

from vivarium.framework.engine import Builder
from vivarium.framework.event import Event
from vivarium.framework.results.exceptions import ResultsConfigurationError
from vivarium.framework.results.observation import BaseObservation
from vivarium.framework.results.observation import Observation
from vivarium.framework.results.stratification import (
Stratification,
get_mapped_col_name,
get_original_col_name,
)
from vivarium.types import ScalarValue
from vivarium.types import ScalarMapper, VectorMapper


class ResultsContext:
Expand Down Expand Up @@ -56,7 +57,9 @@ def __init__(self) -> None:
self.default_stratifications: list[str] = []
self.stratifications: list[Stratification] = []
self.excluded_categories: dict[str, list[str]] = {}
self.observations: defaultdict = defaultdict(lambda: defaultdict(list))
self.observations: defaultdict[
str, defaultdict[tuple[str, tuple[str, ...] | None], list[Observation]]
] = defaultdict(lambda: defaultdict(list))

@property
def name(self) -> str:
Expand Down Expand Up @@ -100,7 +103,7 @@ def add_stratification(
sources: list[str],
categories: list[str],
excluded_categories: list[str] | None,
mapper: Callable[[pd.Series | pd.DataFrame], pd.Series] | Callable[[ScalarValue], str],
mapper: VectorMapper | ScalarMapper | None,
is_vectorized: bool,
) -> None:
"""Add a stratification to the results context.
Expand Down Expand Up @@ -183,11 +186,11 @@ def add_stratification(

def register_observation(
self,
observation_type: Type[BaseObservation],
observation_type: type[Observation],
name: str,
pop_filter: str,
when: str,
**kwargs,
**kwargs: Any,
) -> None:
"""Add an observation to the results context.
Expand Down Expand Up @@ -237,7 +240,15 @@ def register_observation(

def gather_results(
self, population: pd.DataFrame, lifecycle_phase: str, event: Event
) -> Generator[tuple[pd.DataFrame | None, str | None, Callable[[pd.DataFrame, pd.DataFrame], pd.DataFrame] | None], None, None]:
) -> Generator[
tuple[
pd.DataFrame | None,
str | None,
Callable[[pd.DataFrame, pd.DataFrame], pd.DataFrame] | None,
],
None,
None,
]:
"""Generate and yield current results for all observations at this lifecycle
phase and event.
Expand Down Expand Up @@ -290,6 +301,7 @@ def gather_results(
if filtered_pop.empty:
yield None, None, None
else:
pop: pd.DataFrame | DataFrameGroupBy[tuple[str, ...] | str]
if stratification_names is None:
pop = filtered_pop
else:
Expand Down Expand Up @@ -323,7 +335,7 @@ def _filter_population(
@staticmethod
def _get_groups(
stratifications: tuple[str, ...], filtered_pop: pd.DataFrame
) -> DataFrameGroupBy:
) -> DataFrameGroupBy[tuple[str, ...] | str]:
"""Group the population by stratification.
Notes
Expand All @@ -344,7 +356,7 @@ def _get_groups(
)
else:
pop_groups = filtered_pop.groupby(lambda _: "all")
return pop_groups
return pop_groups # type: ignore[return-value]

def _rename_stratification_columns(self, results: pd.DataFrame) -> None:
"""Convert the temporary stratified mapped index names back to their original names."""
Expand Down
3 changes: 2 additions & 1 deletion src/vivarium/framework/results/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

from vivarium.framework.event import Event
from vivarium.framework.results.context import ResultsContext
from vivarium.framework.results.observation import Observation
from vivarium.framework.values import Pipeline
from vivarium.manager import Manager
from vivarium.types import ScalarValue
Expand Down Expand Up @@ -297,7 +298,7 @@ def _bin_data(data: pd.Series | pd.DataFrame) -> pd.Series:

def register_observation(
self,
observation_type,
observation_type: type[Observation],
is_stratified: bool,
name: str,
pop_filter: str,
Expand Down
Loading

0 comments on commit 257fb3d

Please sign in to comment.