diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d7f9861b..a2189974 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,3 +1,7 @@ +**3.2.13 - 12/27/24** + + - Type-hinting: Fix mypy errors in vivarium/examples/boids/ + **3.2.12 - 12/26/24** - Type-hinting: Fix mypy errors in vivarium/framework/engine.py diff --git a/pyproject.toml b/pyproject.toml index 9fd766ee..d1b65d56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,11 +32,6 @@ exclude = [ # You will need to remove the mypy: ignore-errors comment from the file heading as well 'docs/source/conf.py', 'setup.py', - 'src/vivarium/examples/boids/forces.py', - 'src/vivarium/examples/boids/movement.py', - 'src/vivarium/examples/boids/neighbors.py', - 'src/vivarium/examples/boids/population.py', - 'src/vivarium/examples/boids/visualization.py', 'src/vivarium/examples/disease_model/__init__.py', 'src/vivarium/examples/disease_model/disease.py', 'src/vivarium/examples/disease_model/intervention.py', diff --git a/src/vivarium/examples/boids/forces.py b/src/vivarium/examples/boids/forces.py index c960502b..e89674a4 100644 --- a/src/vivarium/examples/boids/forces.py +++ b/src/vivarium/examples/boids/forces.py @@ -1,4 +1,5 @@ -# mypy: ignore-errors + +from __future__ import annotations from abc import ABC, abstractmethod from typing import Any @@ -43,7 +44,7 @@ def setup(self, builder: Builder) -> None: # Pipeline sources and modifiers # ################################## - def apply_force(self, index: pd.Index, acceleration: pd.DataFrame) -> pd.DataFrame: + def apply_force(self, index: pd.Index[int], acceleration: pd.DataFrame) -> pd.DataFrame: neighbors = self.neighbors(index) pop = self.population_view.get(index) pairs = self._get_pairs(neighbors, pop) @@ -56,7 +57,7 @@ def apply_force(self, index: pd.Index, acceleration: pd.DataFrame) -> pd.DataFra max_speed=self.max_speed, ) - acceleration.loc[force.index] += force[["x", "y"]] + acceleration.loc[force.index, ["x", "y"]] += force[["x", "y"]] return acceleration ################## @@ -64,10 +65,10 @@ def apply_force(self, index: pd.Index, acceleration: pd.DataFrame) -> pd.DataFra ################## @abstractmethod - def calculate_force(self, neighbors: pd.DataFrame): + def calculate_force(self, neighbors: pd.DataFrame) -> pd.DataFrame: pass - def _get_pairs(self, neighbors: pd.Series, pop: pd.DataFrame): + def _get_pairs(self, neighbors: pd.Series[int], pop: pd.DataFrame) -> pd.DataFrame: pairs = ( pop.join(neighbors.rename("neighbors")) .reset_index() @@ -91,7 +92,7 @@ def _normalize_and_limit_force( velocity: pd.DataFrame, max_force: float, max_speed: float, - ): + ) -> pd.DataFrame: normalization_factor = np.where( (force.x != 0) | (force.y != 0), max_speed / self._magnitude(force), @@ -111,8 +112,8 @@ def _normalize_and_limit_force( force["y"] *= limit_scaling_factor return force[["x", "y"]] - def _magnitude(self, df: pd.DataFrame): - return np.sqrt(np.square(df.x) + np.square(df.y)) + def _magnitude(self, df: pd.DataFrame) -> pd.Series[float]: + return pd.Series(np.sqrt(np.square(df.x) + np.square(df.y)), dtype=float) class Separation(Force): @@ -125,7 +126,7 @@ class Separation(Force): }, } - def calculate_force(self, neighbors: pd.DataFrame): + def calculate_force(self, neighbors: pd.DataFrame) -> pd.DataFrame: # Push boids apart when they get too close separation_neighbors = neighbors[neighbors.distance < self.config.distance].copy() force_scaling_factor = np.where( @@ -140,17 +141,19 @@ def calculate_force(self, neighbors: pd.DataFrame): separation_neighbors["distance_y"] * force_scaling_factor ) - return ( + force: pd.DataFrame = ( separation_neighbors.groupby("index_self")[["force_x", "force_y"]] .sum() .rename(columns=lambda c: c.replace("force_", "")) ) + return force + class Cohesion(Force): """Push boids together.""" - def calculate_force(self, pairs: pd.DataFrame): + def calculate_force(self, pairs: pd.DataFrame) -> pd.DataFrame: return ( pairs.groupby("index_self")[["distance_x", "distance_y"]] .sum() @@ -161,7 +164,7 @@ def calculate_force(self, pairs: pd.DataFrame): class Alignment(Force): """Push boids toward where others are going.""" - def calculate_force(self, pairs: pd.DataFrame): + def calculate_force(self, pairs: pd.DataFrame) -> pd.DataFrame: return ( pairs.groupby("index_self")[["vx_other", "vy_other"]] .sum() diff --git a/src/vivarium/examples/boids/movement.py b/src/vivarium/examples/boids/movement.py index a1a5cf74..ab589429 100644 --- a/src/vivarium/examples/boids/movement.py +++ b/src/vivarium/examples/boids/movement.py @@ -1,7 +1,8 @@ -# mypy: ignore-errors +from __future__ import annotations import numpy as np import pandas as pd +from vivarium.framework.event import Event from vivarium import Component from vivarium.framework.engine import Builder from vivarium.framework.population import SimulantData @@ -38,7 +39,7 @@ def setup(self, builder: Builder) -> None: # Pipeline sources and modifiers # ################################## - def base_acceleration(self, index: pd.Index) -> pd.DataFrame: + def base_acceleration(self, index: pd.Index[int]) -> pd.DataFrame: return pd.DataFrame(0.0, columns=["x", "y"], index=index) ######################## @@ -59,7 +60,7 @@ def on_initialize_simulants(self, pop_data: SimulantData) -> None: ) self.population_view.update(new_population) - def on_time_step(self, event): + def on_time_step(self, event: Event) -> None: pop = self.population_view.get(event.index) acceleration = self.acceleration(event.index) diff --git a/src/vivarium/examples/boids/neighbors.py b/src/vivarium/examples/boids/neighbors.py index 9004cc11..3cac08ec 100644 --- a/src/vivarium/examples/boids/neighbors.py +++ b/src/vivarium/examples/boids/neighbors.py @@ -1,4 +1,4 @@ -# mypy: ignore-errors +from __future__ import annotations import pandas as pd from scipy import spatial @@ -43,7 +43,7 @@ def on_time_step(self, event: Event) -> None: # Pipeline sources and modifiers # ################################## - def get_neighbors(self, index: pd.Index) -> pd.Series: + def get_neighbors(self, index: pd.Index[int]) -> pd.Series[list[int]]: # type: ignore[type-var] if not self.neighbors_calculated: self._calculate_neighbors() return self._neighbors[index] diff --git a/src/vivarium/examples/boids/population.py b/src/vivarium/examples/boids/population.py index b980298d..ad6abc4a 100644 --- a/src/vivarium/examples/boids/population.py +++ b/src/vivarium/examples/boids/population.py @@ -1,4 +1,3 @@ -# mypy: ignore-errors import numpy as np import pandas as pd diff --git a/src/vivarium/examples/boids/visualization.py b/src/vivarium/examples/boids/visualization.py index d8e1a550..2a41a5a4 100644 --- a/src/vivarium/examples/boids/visualization.py +++ b/src/vivarium/examples/boids/visualization.py @@ -1,34 +1,35 @@ -# mypy: ignore-errors import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation +from vivarium import InteractiveContext -def plot_boids(simulation, plot_velocity=False): + +def plot_boids(simulation: InteractiveContext, plot_velocity: bool=False) -> None: width = simulation.configuration.field.width height = simulation.configuration.field.height pop = simulation.get_population() - plt.figure(figsize=[12, 12]) + plt.figure(figsize=(12, 12)) plt.scatter(pop.x, pop.y, color=pop.color) if plot_velocity: plt.quiver(pop.x, pop.y, pop.vx, pop.vy, color=pop.color, width=0.002) plt.xlabel("x") plt.ylabel("y") - plt.axis([0, width, 0, height]) + plt.axis((0, width, 0, height)) plt.show() -def plot_boids_animated(simulation): +def plot_boids_animated(simulation: InteractiveContext) -> FuncAnimation: width = simulation.configuration.field.width height = simulation.configuration.field.height pop = simulation.get_population() - fig = plt.figure(figsize=[12, 12]) + fig = plt.figure(figsize=(12, 12)) ax = fig.add_subplot(111) s = ax.scatter(pop.x, pop.y, color=pop.color) plt.xlabel("x") plt.ylabel("y") - plt.axis([0, width, 0, height]) + plt.axis((0, width, 0, height)) frames = range(2_000) frame_pops = [] @@ -36,7 +37,7 @@ def plot_boids_animated(simulation): simulation.step() frame_pops.append(simulation.get_population()[["x", "y"]]) - def animate(i): + def animate(i: int) -> None: s.set_offsets(frame_pops[i]) - return FuncAnimation(fig, animate, frames=frames, interval=10) + return FuncAnimation(fig, animate, frames=frames, interval=10) # type: ignore[arg-type]