From 6379644bf0c6bd574769d1eadd73a318bf964144 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Mon, 14 Oct 2024 14:35:04 +0200 Subject: [PATCH] refactor: memory usage (#1238 #1276 #1277) --- openfisca_core/holders/__init__.py | 2 -- openfisca_core/holders/holder.py | 5 ++-- openfisca_core/holders/memory_usage.py | 26 ------------------- openfisca_core/holders/types.py | 4 +-- .../populations/_core_population.py | 6 ++--- openfisca_core/populations/population.py | 14 +++++----- openfisca_core/populations/types.py | 8 ++++++ openfisca_core/types.py | 12 ++++++++- 8 files changed, 33 insertions(+), 44 deletions(-) delete mode 100644 openfisca_core/holders/memory_usage.py diff --git a/openfisca_core/holders/__init__.py b/openfisca_core/holders/__init__.py index 127d7fa92..a120a671b 100644 --- a/openfisca_core/holders/__init__.py +++ b/openfisca_core/holders/__init__.py @@ -24,11 +24,9 @@ from . import types from .helpers import set_input_dispatch_by_period, set_input_divide_by_period from .holder import Holder -from .memory_usage import MemoryUsage __all__ = [ "Holder", - "MemoryUsage", "set_input_dispatch_by_period", "set_input_divide_by_period", "types", diff --git a/openfisca_core/holders/holder.py b/openfisca_core/holders/holder.py index e79ef9ad2..4bac41b66 100644 --- a/openfisca_core/holders/holder.py +++ b/openfisca_core/holders/holder.py @@ -19,7 +19,6 @@ ) from . import types as t -from .memory_usage import MemoryUsage class Holder: @@ -95,7 +94,7 @@ def get_array(self, period): return self._disk_storage.get(period) return None - def get_memory_usage(self) -> MemoryUsage: + def get_memory_usage(self) -> t.MemoryUsage: """Get data about the virtual memory usage of the Holder. Returns: @@ -136,7 +135,7 @@ def get_memory_usage(self) -> MemoryUsage: 'total_nb_bytes': 0... """ - usage = MemoryUsage( + usage = t.MemoryUsage( nb_cells_by_array=self.population.count, dtype=self.variable.dtype, ) diff --git a/openfisca_core/holders/memory_usage.py b/openfisca_core/holders/memory_usage.py deleted file mode 100644 index 2d344318e..000000000 --- a/openfisca_core/holders/memory_usage.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing_extensions import TypedDict - -import numpy - - -class MemoryUsage(TypedDict, total=False): - """Virtual memory usage of a Holder. - - Attributes: - cell_size: The amount of bytes assigned to each value. - dtype: The :mod:`numpy.dtype` of any, each, and every value. - nb_arrays: The number of periods for which the Holder contains values. - nb_cells_by_array: The number of entities in the current Simulation. - nb_requests: The number of times the Variable has been computed. - nb_requests_by_array: Average times a stored array has been read. - total_nb_bytes: The total number of bytes used by the Holder. - - """ - - cell_size: int - dtype: numpy.dtype - nb_arrays: int - nb_cells_by_array: int - nb_requests: int - nb_requests_by_array: int - total_nb_bytes: int diff --git a/openfisca_core/holders/types.py b/openfisca_core/holders/types.py index 355fe71e4..7137b8648 100644 --- a/openfisca_core/holders/types.py +++ b/openfisca_core/holders/types.py @@ -1,3 +1,3 @@ -from openfisca_core.types import CorePopulation, Holder +from openfisca_core.types import CorePopulation, Holder, MemoryUsage -__all__ = ["CorePopulation", "Holder"] +__all__ = ["CorePopulation", "Holder", "MemoryUsage"] diff --git a/openfisca_core/populations/_core_population.py b/openfisca_core/populations/_core_population.py index 908fd3654..ab97be383 100644 --- a/openfisca_core/populations/_core_population.py +++ b/openfisca_core/populations/_core_population.py @@ -27,11 +27,11 @@ class Option(strenum.StrEnum): class Calculate(NamedTuple): variable: str period: t.Period - option: Sequence[str] | None + option: None | Sequence[str] class MemoryUsageByVariable(TypedDict, total=False): - by_variable: dict[str, holders.MemoryUsage] + by_variable: dict[str, t.MemoryUsage] total_nb_bytes: int @@ -186,7 +186,7 @@ def check_array_compatible_with_entity(self, array: t.FloatArray) -> None: def check_period_validity( self, variable_name: str, - period: int | str | Period | None, + period: None | t.PeriodLike, ) -> None: if isinstance(period, (int, str, periods.Period)): return diff --git a/openfisca_core/populations/population.py b/openfisca_core/populations/population.py index 0f5275cfa..24742ab0a 100644 --- a/openfisca_core/populations/population.py +++ b/openfisca_core/populations/population.py @@ -12,7 +12,7 @@ class Population(CorePopulation): def __init__(self, entity: t.SingleEntity) -> None: super().__init__(entity) - def clone(self, simulation: Simulation) -> t.CorePopulation: + def clone(self, simulation: t.Simulation) -> t.CorePopulation: result = Population(self.entity) result.simulation = simulation result._holders = { @@ -38,7 +38,7 @@ def __getattr__(self, attribute: str) -> projectors.Projector: # Helpers @projectors.projectable - def has_role(self, role: Role) -> Array[bool] | None: + def has_role(self, role: t.Role) -> None | t.BoolArray: """Check if a person has a given role within its `GroupEntity`. Example: @@ -63,10 +63,10 @@ def has_role(self, role: Role) -> Array[bool] | None: @projectors.projectable def value_from_partner( self, - array: Array[float], + array: t.FloatArray, entity: projectors.Projector, - role: Role, - ) -> Array[float] | None: + role: t.Role, + ) -> None | t.FloatArray: self.check_array_compatible_with_entity(array) self.entity.check_role_validity(role) @@ -89,9 +89,9 @@ def value_from_partner( def get_rank( self, entity: Population, - criteria: Array[float], + criteria: t.FloatArray, condition: bool = True, - ) -> Array[int]: + ) -> t.IntArray: """Get the rank of a person within an entity according to a criteria. The person with rank 0 has the minimum value of criteria. If condition is specified, then the persons who don't respect it are not taken into account and their rank is -1. diff --git a/openfisca_core/populations/types.py b/openfisca_core/populations/types.py index c0056497d..684c5456c 100644 --- a/openfisca_core/populations/types.py +++ b/openfisca_core/populations/types.py @@ -12,8 +12,10 @@ EntityKey, GroupEntity, Holder, + MemoryUsage, Period, PeriodStr, + Role, Simulation, SingleEntity, SinglePopulation, @@ -24,11 +26,15 @@ bool_ as BoolDType, float32 as FloatDType, generic as VarDType, + int32 as IntDType, str_ as StrDType, ) # Commons +#: Type alias for an array of strings. +IntArray: TypeAlias = Array[IntDType] + #: Type alias for an array of strings. StrArray: TypeAlias = Array[StrDType] @@ -63,7 +69,9 @@ "EntityKey", "GroupEntity", "Holder", + "MemoryUsage", "Period", + "Role", "Simulation", "SingleEntity", "SinglePopulation", diff --git a/openfisca_core/types.py b/openfisca_core/types.py index 04e2a7abb..cdf0606d9 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -3,7 +3,7 @@ from collections.abc import Iterable, Mapping, Sequence, Sized from numpy.typing import DTypeLike, NDArray from typing import NewType, TypeVar, Union -from typing_extensions import Protocol, Self, TypeAlias +from typing_extensions import Protocol, Self, TypeAlias, TypedDict import abc import enum @@ -136,6 +136,16 @@ def clone(self, population: CorePopulation, /) -> Holder: ... def get_memory_usage(self, /) -> Mapping[str, object]: ... +class MemoryUsage(TypedDict, total=False): + cell_size: int + dtype: DTypeLike + nb_arrays: int + nb_cells_by_array: int + nb_requests: int + nb_requests_by_array: int + total_nb_bytes: int + + # Parameters