From 93c161c102e456b928067bae64deb3b2aef70e0c Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 6 Aug 2024 15:10:14 +0900 Subject: [PATCH 01/89] target version of develop branch is changed --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 898b77f..ba868c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "phlower" -version = "0.0.1" +version = "0.2.0" description = "This is a Python package which helps you handle GNN especially for physics problems." authors = ["sakamoto "] readme = "README.md" From 7cca6ba125ac2fc45d3e04b9c0afd0f01b1530e1 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 7 Aug 2024 10:33:54 +0900 Subject: [PATCH 02/89] add progress bar while batch iteration --- src/phlower/services/trainer/_trainer.py | 12 ++++++++++++ src/phlower/utils/_progress_bar.py | 5 +++-- tests/e2e_tests/test_train.py | 2 -- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index b4f2ed6..4f65252 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -111,6 +111,9 @@ def train( tqdm.write(record_io.get_header()) self._timer.start() + _train_batch_pbar = PhlowerProgressBar(total=len(train_dataset)) + _val_batch_pbar = PhlowerProgressBar(total=len(validation_dataset)) + for epoch in range(self._trainer_setting.n_epoch): train_losses: list[float] = [] validation_losses: list[float] = [] @@ -132,6 +135,11 @@ def train( loss.backward() self._optimizer.step() + _train_batch_pbar.update( + trick=self._trainer_setting.batch_size, + desc=f"batch train loss: {train_losses[-1]:.3f}", + ) + self._model.eval() for val_batch in validation_loader: with torch.no_grad(): @@ -147,6 +155,10 @@ def train( validation_losses.append( val_loss.detach().to_tensor().float().item() ) + _val_batch_pbar.update( + trick=self._trainer_setting.batch_size, + desc=f"batch val loss: {validation_losses[-1]}" + ) train_loss = np.average(train_losses) validation_loss = np.average(validation_losses) diff --git a/src/phlower/utils/_progress_bar.py b/src/phlower/utils/_progress_bar.py index 33587c1..54c4772 100644 --- a/src/phlower/utils/_progress_bar.py +++ b/src/phlower/utils/_progress_bar.py @@ -45,5 +45,6 @@ def update(self, trick: int, *, desc: str = None) -> None: def destroy(self) -> None: # Unlike tqdm.reset, this method does not show next progress bar. - self._pbar.close() - self._pbar = None + if self._pbar is not None: + self._pbar.close() + self._pbar = None diff --git a/tests/e2e_tests/test_train.py b/tests/e2e_tests/test_train.py index 22b88e9..a94ac5b 100644 --- a/tests/e2e_tests/test_train.py +++ b/tests/e2e_tests/test_train.py @@ -113,8 +113,6 @@ def test__training_with_multiple_batch_size(prepare_sample_preprocessed_files): def test__simple_training(simple_training): loss: PhlowerTensor = simple_training - print(loss) - print(loss.shape) assert loss.has_dimension assert not torch.isinf(loss.to_tensor()) assert not torch.isnan(loss.to_tensor()) From e21a41d810bf75e1d4e8e9402e8de9d4033e64c9 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 7 Aug 2024 10:35:44 +0900 Subject: [PATCH 03/89] fix lint warnings --- src/phlower/services/trainer/_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index 4f65252..4abfec6 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -157,7 +157,7 @@ def train( ) _val_batch_pbar.update( trick=self._trainer_setting.batch_size, - desc=f"batch val loss: {validation_losses[-1]}" + desc=f"batch val loss: {validation_losses[-1]}", ) train_loss = np.average(train_losses) From e8a7a2ec16361553ae7cc28ade49d86583e376cc Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 7 Aug 2024 16:12:36 +0900 Subject: [PATCH 04/89] add interfact to handle reference from one module to another module. Ex. Share module --- src/phlower/nn/__init__.py | 2 + src/phlower/nn/_core_modules/__init__.py | 3 +- src/phlower/nn/_core_modules/_concatenator.py | 16 +++- src/phlower/nn/_core_modules/_gcn.py | 14 +++- src/phlower/nn/_core_modules/_mlp.py | 14 +++- src/phlower/nn/_core_modules/_share.py | 77 +++++++++++++++++++ src/phlower/nn/_group_module.py | 23 ++++-- src/phlower/nn/_interface_module.py | 23 +++++- src/phlower/nn/_phlower_module_adpter.py | 9 ++- src/phlower/settings/_group_settings.py | 69 ++++++++++------- src/phlower/settings/_interface.py | 44 +++++++++++ .../settings/_module_settings/__init__.py | 12 +-- .../_module_settings/_concatenator_setting.py | 15 +++- .../settings/_module_settings/_gcn_setting.py | 15 +++- .../settings/_module_settings/_interface.py | 13 ---- .../settings/_module_settings/_mlp_setting.py | 15 +++- .../_module_settings/_share_setting.py | 50 ++++++++++++ src/phlower/utils/exceptions.py | 3 + tests/samples/settings/with_share_nn.yml | 46 +++++++++++ tests/test_nn/test_core_modules/test_share.py | 50 ++++++++++++ tests/test_nn/test_group_module.py | 18 +++++ .../share_settings/check_gcn_share_nodes.yml | 50 ++++++++++++ .../share_settings/check_mlp_share_nodes.yml | 50 ++++++++++++ .../data/share_settings/with_share_nn.yml | 46 +++++++++++ .../test_share_settings.py | 30 ++++++++ 25 files changed, 639 insertions(+), 68 deletions(-) create mode 100644 src/phlower/nn/_core_modules/_share.py create mode 100644 src/phlower/settings/_interface.py delete mode 100644 src/phlower/settings/_module_settings/_interface.py create mode 100644 src/phlower/settings/_module_settings/_share_setting.py create mode 100644 tests/samples/settings/with_share_nn.yml create mode 100644 tests/test_nn/test_core_modules/test_share.py create mode 100644 tests/test_nn/test_group_module.py create mode 100644 tests/test_settings/test_module_settings/data/share_settings/check_gcn_share_nodes.yml create mode 100644 tests/test_settings/test_module_settings/data/share_settings/check_mlp_share_nodes.yml create mode 100644 tests/test_settings/test_module_settings/data/share_settings/with_share_nn.yml create mode 100644 tests/test_settings/test_module_settings/test_share_settings.py diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 7844fd3..1af9012 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -1,4 +1,6 @@ from phlower.nn._core_modules._concatenator import Concatenator from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._share import Share from phlower.nn._group_module import PhlowerGroupModule +from phlower.nn._interface_module import IPhlowerCoreModule diff --git a/src/phlower/nn/_core_modules/__init__.py b/src/phlower/nn/_core_modules/__init__.py index 51d8c94..17800d6 100644 --- a/src/phlower/nn/_core_modules/__init__.py +++ b/src/phlower/nn/_core_modules/__init__.py @@ -1,9 +1,10 @@ from phlower.nn._core_modules._concatenator import Concatenator from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._share import Share from phlower.nn._interface_module import IPhlowerCoreModule -_all_models: list[IPhlowerCoreModule] = [GCN, MLP, Concatenator] +_all_models: list[IPhlowerCoreModule] = [GCN, MLP, Concatenator, Share] _name2model = {cls.get_nn_name(): cls for cls in _all_models} diff --git a/src/phlower/nn/_core_modules/_concatenator.py b/src/phlower/nn/_core_modules/_concatenator.py index 5a0288c..2dbbfae 100644 --- a/src/phlower/nn/_core_modules/_concatenator.py +++ b/src/phlower/nn/_core_modules/_concatenator.py @@ -6,11 +6,12 @@ from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils -from phlower.nn._interface_module import IPhlowerCoreModule +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) from phlower.settings._module_settings import ConcatenatorSetting -# from phlower.settings - class Concatenator(IPhlowerCoreModule, torch.nn.Module): """Concatenator""" @@ -36,17 +37,26 @@ def get_nn_name(cls) -> str: """ return "Concatenator" + @classmethod + def need_resolve(cls) -> bool: + return False + def __init__(self, activation: str, nodes: list[int] = None): super().__init__() self._nodes = nodes self._activation_name = activation self._activation_func = _utils.ActivationSelector.select(activation) + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + def forward( self, data: IPhlowerTensorCollections, *, supports: dict[str, PhlowerTensor] | None = None, + **kwards, ) -> PhlowerTensor: """forward function which overloads torch.nn.Module diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index d12a546..6061e3d 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -5,7 +5,10 @@ from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils -from phlower.nn._interface_module import IPhlowerCoreModule +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) from phlower.settings._module_settings import GCNSetting @@ -33,6 +36,10 @@ def get_nn_name(cls) -> str: """ return "GCN" + @classmethod + def need_resolve(cls) -> bool: + return False + def __init__( self, nodes: list[int], @@ -59,11 +66,16 @@ def __init__( self._repeat = repeat self._factor = factor + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + def forward( self, data: IPhlowerTensorCollections, *, supports: dict[str, PhlowerTensor], + **kwards, ) -> PhlowerTensor: """forward function which overload torch.nn.Module diff --git a/src/phlower/nn/_core_modules/_mlp.py b/src/phlower/nn/_core_modules/_mlp.py index 29fe37e..bd0074c 100644 --- a/src/phlower/nn/_core_modules/_mlp.py +++ b/src/phlower/nn/_core_modules/_mlp.py @@ -6,7 +6,10 @@ from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils -from phlower.nn._interface_module import IPhlowerCoreModule +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) from phlower.settings._module_settings import MLPSetting @@ -34,6 +37,10 @@ def get_nn_name(cls) -> str: """ return "MLP" + @classmethod + def need_resolve(cls) -> bool: + return False + def __init__( self, nodes: list[int], @@ -54,11 +61,16 @@ def __init__( self._nodes = nodes self._activations = activations + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + def forward( self, data: IPhlowerTensorCollections, *, supports: dict[str, PhlowerTensor] | None = None, + **kwards, ) -> PhlowerTensor: """forward function which overloads torch.nn.Module diff --git a/src/phlower/nn/_core_modules/_share.py b/src/phlower/nn/_core_modules/_share.py new file mode 100644 index 0000000..131ec4b --- /dev/null +++ b/src/phlower/nn/_core_modules/_share.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import torch + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings import ShareSetting +from phlower.utils.exceptions import NotFoundReferenceModuleError + + +class Share(IPhlowerCoreModule, torch.nn.Module): + """ + Share module is a reference to another module. + Share module itself does not have any trainable parameters. + """ + + @classmethod + def from_setting(cls, setting: ShareSetting) -> Share: + return Share(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + return "Share" + + @classmethod + def need_resolve(cls) -> bool: + return True + + def __init__(self, reference_name: str, **kwards) -> None: + super().__init__() + + self._reference_name = reference_name + self._reference: IPhlowerCoreModule | None = None + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: + assert parent is not None + + try: + self._reference = parent.search_module(self._reference_name) + except KeyError as ex: + raise NotFoundReferenceModuleError( + f"Reference module {self._reference_name} is not found " + "in the same group." + ) from ex + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + **kwards, + ) -> PhlowerTensor: + """forward function which overload torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor]): + sparse tensor objects + + Returns: + PhlowerTensor: + Tensor object + """ + if self._reference is None: + raise ValueError( + "reference module in Share module is not set. " + "Please check that `resolve` function is called." + ) + + return self._reference.forward(data, supports=supports, **kwards) diff --git a/src/phlower/nn/_group_module.py b/src/phlower/nn/_group_module.py index a94782c..a8907e8 100644 --- a/src/phlower/nn/_group_module.py +++ b/src/phlower/nn/_group_module.py @@ -13,14 +13,19 @@ reduce_collections, ) from phlower.io._files import IPhlowerCheckpointFile -from phlower.nn._interface_module import IPhlowerModuleAdapter +from phlower.nn._interface_module import ( + IPhlowerModuleAdapter, + IReadonlyReferenceGroup, +) from phlower.nn._phlower_module_adpter import PhlowerModuleAdapter from phlower.services.drawers import MermaidDrawer from phlower.settings._group_settings import GroupModuleSetting, ModuleSetting from phlower.utils.enums import TrainerSavedKeyType -class PhlowerGroupModule(IPhlowerModuleAdapter, torch.nn.Module): +class PhlowerGroupModule( + IPhlowerModuleAdapter, IReadonlyReferenceGroup, torch.nn.Module +): @classmethod def from_setting(cls, setting: GroupModuleSetting) -> Self: _modules: list[IPhlowerModuleAdapter] = [] @@ -81,7 +86,7 @@ def get_n_nodes(self) -> list[int]: def resolve(self): for module in self._phlower_modules: - module.resolve() + module.resolve(parent=self) # topological-sort stream = dagstream.DagStream() @@ -122,7 +127,8 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor], + supports: dict[str, PhlowerTensor] | None = None, + **kwards, ) -> IPhlowerTensorCollections: results = phlower_tensor_collection({}) @@ -140,7 +146,7 @@ def forward( args = reduce_collections(node.get_received_args()) _module: PhlowerModuleAdapter = node.get_user_function() - _result = _module.forward(args, supports=supports) + _result = _module.forward(args, supports=supports, **kwards) dag_modules.send(node.mut_name, _result) dag_modules.done(node.mut_name) @@ -153,6 +159,13 @@ def forward( def get_destinations(self) -> list[str]: return self._destinations + def search_module(self, name: str) -> IPhlowerModuleAdapter: + for _module in self._phlower_modules: + if _module.name == name: + return _module + + raise KeyError(f"Module {name} is not found in group {self.name}.") + def load_checkpoint_file( self, checkpoint_file: IPhlowerCheckpointFile, diff --git a/src/phlower/nn/_interface_module.py b/src/phlower/nn/_interface_module.py index 260f1fd..e91a5fe 100644 --- a/src/phlower/nn/_interface_module.py +++ b/src/phlower/nn/_interface_module.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import abc import pathlib @@ -17,6 +19,15 @@ def from_setting(cls, setting: IPhlowerLayerParameters) -> Self: ... @abc.abstractmethod def get_nn_name(cls) -> str: ... + @classmethod + @abc.abstractmethod + def need_resolve(cls) -> bool: ... + + @abc.abstractmethod + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + @abc.abstractmethod def forward( self, @@ -26,20 +37,28 @@ def forward( ) -> PhlowerTensor: ... +class IReadonlyReferenceGroup(metaclass=abc.ABCMeta): + @abc.abstractmethod + def search_module(self, name: str) -> IPhlowerModuleAdapter: ... + + class IPhlowerModuleAdapter(metaclass=abc.ABCMeta): @property @abc.abstractmethod def name(self) -> str: ... @abc.abstractmethod - def resolve(self) -> None: ... + def resolve( + cls, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... @abc.abstractmethod def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor], + supports: dict[str, PhlowerTensor] | None = None, + **kwards, ) -> IPhlowerTensorCollections: ... @abc.abstractmethod diff --git a/src/phlower/nn/_phlower_module_adpter.py b/src/phlower/nn/_phlower_module_adpter.py index fdb1340..979b87c 100644 --- a/src/phlower/nn/_phlower_module_adpter.py +++ b/src/phlower/nn/_phlower_module_adpter.py @@ -13,6 +13,7 @@ from phlower.nn._interface_module import ( IPhlowerCoreModule, IPhlowerModuleAdapter, + IReadonlyReferenceGroup, ) from phlower.settings._group_settings import ModuleSetting @@ -57,7 +58,13 @@ def __init__( def get_destinations(self) -> list[str]: return self._destinations - def resolve(self) -> None: ... + def resolve( + self, *, parent: IReadonlyReferenceGroup | None, **kwards + ) -> None: + if not self._layer.need_resolve(): + return + + self._layer.resolve(parent=parent, **kwards) def get_n_nodes(self) -> list[int]: return self._n_nodes diff --git a/src/phlower/settings/_group_settings.py b/src/phlower/settings/_group_settings.py index d1806dd..f4228ff 100644 --- a/src/phlower/settings/_group_settings.py +++ b/src/phlower/settings/_group_settings.py @@ -1,6 +1,5 @@ from __future__ import annotations -import abc from typing import Literal import dagstream @@ -10,8 +9,12 @@ from pydantic import dataclasses as dc from typing_extensions import Self +from phlower.settings._interface import ( + IModuleSetting, + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) from phlower.settings._module_parameter_setting import PhlowerModuleParameters -from phlower.settings._module_settings import gather_input_dims from phlower.utils.exceptions import ( PhlowerModuleCycleError, PhlowerModuleDuplicateKeyError, @@ -21,27 +24,15 @@ ) -class IModuleSetting(metaclass=abc.ABCMeta): - @abc.abstractmethod - def get_name(self) -> str: ... - - @abc.abstractmethod - def get_destinations(self) -> list[str]: ... - - @abc.abstractmethod - def get_output_info(self) -> dict[str, int]: ... - - @abc.abstractmethod - def resolve(self, *resolved_outputs: dict[str, int]) -> None: ... - - @dc.dataclass(frozen=True, config=pydantic.ConfigDict(extra="forbid")) class ModuleIOSetting: name: str n_dim: int -class GroupModuleSetting(IModuleSetting, pydantic.BaseModel): +class GroupModuleSetting( + IModuleSetting, IReadOnlyReferenceGroupSetting, pydantic.BaseModel +): name: str """ name of group @@ -129,8 +120,20 @@ def get_output_keys(self) -> list[str]: def get_output_info(self) -> dict[str, int]: return {output.name: output.n_dim for output in self.outputs} + def search_module_setting(self, name: str) -> IPhlowerLayerParameters: + for module in self.modules: + if module.name == name: + return module.nn_parameters + + raise KeyError( + f"ModuleSetting {name} is not found in GroupSetting {self.name}." + ) + def resolve( - self, *resolved_outputs: dict[str, int], is_first: bool = False + self, + *resolved_outputs: dict[str, int], + is_first: bool = False, + **kwards, ) -> None: if not is_first: self._check_keys(*resolved_outputs) @@ -139,7 +142,7 @@ def resolve( input_info = {v.name: v.n_dim for v in self.inputs} try: - results = _resolve_modules(input_info, self.modules) + results = _resolve_modules(input_info, self.modules, self) except DagStreamCycleError as ex: raise PhlowerModuleCycleError( f"A cycle is detected in {self.name}." @@ -273,8 +276,16 @@ def get_destinations(self) -> list[str]: def get_output_info(self) -> dict[str, int]: return {self.output_key: self.nn_parameters.get_n_nodes()[-1]} - def resolve(self, *resolved_outputs: dict[str, int]) -> None: + def resolve( + self, + *resolved_outputs: dict[str, int], + parent: IReadOnlyReferenceGroupSetting | None = None, + ) -> None: self._check_keys(*resolved_outputs) + + if self.nn_parameters.need_reference: + self.nn_parameters.get_reference(parent) + _resolved_nodes = self._resolve_nodes(*resolved_outputs) # NOTE: overwrite nodes self.nn_parameters.overwrite_nodes(_resolved_nodes) @@ -306,8 +317,8 @@ def _resolve_nodes(self, *resolved_outputs: dict[str, int]) -> list[int]: key2node.update(output) try: - first_node = gather_input_dims( - self.nn_type, *(key2node[key] for key in self.input_keys) + first_node = self.nn_parameters.gather_input_dims( + *(key2node[key] for key in self.input_keys) ) except ValueError as ex: raise PhlowerModuleNodeDimSizeError( @@ -345,13 +356,19 @@ def __init__(self, setting: IModuleSetting) -> None: def name(self) -> str: return self._setting.get_name() - def __call__(self, *resolved_output: dict[str, int]) -> int: - self._setting.resolve(*resolved_output) + def __call__( + self, + *resolved_output: dict[str, int], + parent: IReadOnlyReferenceGroupSetting | None = None, + ) -> int: + self._setting.resolve(*resolved_output, parent=parent) return self._setting.get_output_info() def _resolve_modules( - starts: dict[str, int], modules: list[IModuleSetting] + starts: dict[str, int], + modules: list[IModuleSetting], + parent: IReadOnlyReferenceGroupSetting | None = None, ) -> list[dict[str, int]]: stream = dagstream.DagStream() resolvers = [_SettingResolverAdapter(layer) for layer in modules] @@ -367,6 +384,6 @@ def _resolve_modules( _dag = stream.construct() executor = dagstream.StreamExecutor(_dag) - results = executor.run(first_args=(starts,)) + results = executor.run(parent=parent, first_args=(starts,)) return list(results.values()) diff --git a/src/phlower/settings/_interface.py b/src/phlower/settings/_interface.py new file mode 100644 index 0000000..23acbe8 --- /dev/null +++ b/src/phlower/settings/_interface.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import abc + + +class IReadOnlyReferenceGroupSetting(metaclass=abc.ABCMeta): + @abc.abstractmethod + def search_module_setting(self, name: str) -> IPhlowerLayerParameters: ... + + +class IModuleSetting(metaclass=abc.ABCMeta): + @abc.abstractmethod + def get_name(self) -> str: ... + + @abc.abstractmethod + def get_destinations(self) -> list[str]: ... + + @abc.abstractmethod + def get_output_info(self) -> dict[str, int]: ... + + @abc.abstractmethod + def resolve( + self, + *resolved_outputs: dict[str, int], + parent: IReadOnlyReferenceGroupSetting | None = None, + ) -> None: ... + + +class IPhlowerLayerParameters(metaclass=abc.ABCMeta): + @abc.abstractmethod + def gather_input_dims(self, *input_dims: int) -> int: ... + + @abc.abstractmethod + def get_n_nodes(self) -> list[int] | None: ... + + @abc.abstractmethod + def overwrite_nodes(self, nodes: list[int]) -> None: ... + + @property + @abc.abstractmethod + def need_reference(self) -> bool: ... + + @abc.abstractmethod + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): ... diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index 4eca214..d576583 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -1,21 +1,21 @@ +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) from phlower.settings._module_settings._concatenator_setting import ( ConcatenatorSetting, ) from phlower.settings._module_settings._gcn_setting import GCNSetting -from phlower.settings._module_settings._interface import IPhlowerLayerParameters from phlower.settings._module_settings._mlp_setting import MLPSetting +from phlower.settings._module_settings._share_setting import ShareSetting _name_to_setting: dict[str, IPhlowerLayerParameters] = { "GCN": GCNSetting, "MLP": MLPSetting, "Concatenator": ConcatenatorSetting, + "Share": ShareSetting, } def check_exist_module(name: str) -> bool: return name in _name_to_setting - - -def gather_input_dims(name: str, *input_dims: int): - setting = _name_to_setting[name] - return setting.gather_input_dims(*input_dims) diff --git a/src/phlower/settings/_module_settings/_concatenator_setting.py b/src/phlower/settings/_module_settings/_concatenator_setting.py index 973adb4..341e2a1 100644 --- a/src/phlower/settings/_module_settings/_concatenator_setting.py +++ b/src/phlower/settings/_module_settings/_concatenator_setting.py @@ -3,7 +3,10 @@ import pydantic from pydantic import Field -from ._interface import IPhlowerLayerParameters +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) class ConcatenatorSetting(IPhlowerLayerParameters, pydantic.BaseModel): @@ -15,8 +18,7 @@ class ConcatenatorSetting(IPhlowerLayerParameters, pydantic.BaseModel): # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid") - @classmethod - def gather_input_dims(cls, *input_dims: int) -> int: + def gather_input_dims(self, *input_dims: int) -> int: assert len(input_dims) > 0 sum_dim = sum(v for v in input_dims) return sum_dim @@ -34,6 +36,13 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: ) return vals + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return + def get_n_nodes(self) -> list[int] | None: return self.nodes diff --git a/src/phlower/settings/_module_settings/_gcn_setting.py b/src/phlower/settings/_module_settings/_gcn_setting.py index f96c7ac..6bb204a 100644 --- a/src/phlower/settings/_module_settings/_gcn_setting.py +++ b/src/phlower/settings/_module_settings/_gcn_setting.py @@ -4,7 +4,10 @@ from pydantic import Field from typing_extensions import Self -from ._interface import IPhlowerLayerParameters +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) class GCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): @@ -21,8 +24,7 @@ class GCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid") - @classmethod - def gather_input_dims(cls, *input_dims: int) -> int: + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError("only one input is allowed in GCN.") return input_dims[0] @@ -65,3 +67,10 @@ def get_n_nodes(self) -> list[int]: def overwrite_nodes(self, nodes: list[int]) -> None: self.nodes = nodes + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return diff --git a/src/phlower/settings/_module_settings/_interface.py b/src/phlower/settings/_module_settings/_interface.py deleted file mode 100644 index 185e0ee..0000000 --- a/src/phlower/settings/_module_settings/_interface.py +++ /dev/null @@ -1,13 +0,0 @@ -import abc - - -class IPhlowerLayerParameters(metaclass=abc.ABCMeta): - @classmethod - @abc.abstractmethod - def gather_input_dims(cls, *input_dims: int) -> int: ... - - @abc.abstractmethod - def get_n_nodes(self) -> list[int] | None: ... - - @abc.abstractmethod - def overwrite_nodes(self, nodes: list[int]) -> None: ... diff --git a/src/phlower/settings/_module_settings/_mlp_setting.py b/src/phlower/settings/_module_settings/_mlp_setting.py index deeba9f..b4d02ce 100644 --- a/src/phlower/settings/_module_settings/_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_mlp_setting.py @@ -4,7 +4,10 @@ from pydantic import Field from typing_extensions import Self -from ._interface import IPhlowerLayerParameters +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) class MLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): @@ -15,8 +18,7 @@ class MLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) bias: bool = Field(False, frozen=True) - @classmethod - def gather_input_dims(cls, *input_dims: int) -> int: + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError("only one input is allowed in GCN.") return input_dims[0] @@ -59,3 +61,10 @@ def get_n_nodes(self) -> list[int] | None: def overwrite_nodes(self, nodes: list[int]) -> None: self.nodes = nodes + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return diff --git a/src/phlower/settings/_module_settings/_share_setting.py b/src/phlower/settings/_module_settings/_share_setting.py new file mode 100644 index 0000000..808f40d --- /dev/null +++ b/src/phlower/settings/_module_settings/_share_setting.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import pydantic +from pydantic import Field + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) +from phlower.utils.exceptions import NotFoundReferenceModuleError + + +class ShareSetting(IPhlowerLayerParameters, pydantic.BaseModel): + reference_name: str = Field(..., frozen=True) + reference: IPhlowerLayerParameters | None = Field(None, exclude=True) + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict( + extra="forbid", arbitrary_types_allowed=True + ) + + def gather_input_dims(self, *input_dims: int) -> int: + return self.reference.gather_input_dims(*input_dims) + + def get_n_nodes(self) -> list[int] | None: + return self.reference.get_n_nodes() + + def check_exist_reference(self) -> None: + if self.reference is not None: + return + + raise ValueError( + f"Reference setting {self.reference_name} in Share Module is None." + "Please check that `get_reference` method has been called." + ) + + def overwrite_nodes(self, nodes: list[int]) -> None: ... + + @property + def need_reference(self) -> bool: + return True + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + try: + self.reference = parent.search_module_setting(self.reference_name) + except KeyError as ex: + raise NotFoundReferenceModuleError( + f"Reference module {self.reference_name} is not found " + "in the same group." + ) from ex diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index fde6954..e7aa38f 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -41,3 +41,6 @@ class PhlowerFeatureStoreOverwriteError(ValueError): """This error raises when to try to overwrite feature store item""" ... + + +class NotFoundReferenceModuleError(ValueError): ... diff --git a/tests/samples/settings/with_share_nn.yml b/tests/samples/settings/with_share_nn.yml new file mode 100644 index 0000000..09f5f6f --- /dev/null +++ b/tests/samples/settings/with_share_nn.yml @@ -0,0 +1,46 @@ + +model: + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - Share0 + nn_parameters: + nodes: [-1, 20, 10] + activations: ["identity", "identity"] + + - nn_type: Share + name: Share0 + input_keys: + - mlp0 + output_key: share1 + destinations: + - GCN0 + nn_parameters: + reference_name: MLP0 + + - nn_type: GCN + name: GCN0 + input_keys: + - share1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_nn/test_core_modules/test_share.py b/tests/test_nn/test_core_modules/test_share.py new file mode 100644 index 0000000..20c6d26 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_share.py @@ -0,0 +1,50 @@ +from unittest import mock + +import numpy as np +import pytest +import torch + +from phlower import PhlowerTensor +from phlower.collections import phlower_tensor_collection +from phlower.nn import MLP, Share +from phlower.nn._interface_module import IReadonlyReferenceGroup + + +def test__can_call_parameters(): + model = Share(reference_name="MLP0") + MLP0 = MLP(nodes=[10, 10]) + model._reference = MLP0 + + # To check Concatenator inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "mlp_nodes", + [([10, 10]), ([20, 10, 100])], +) +def test__reference_same_object(mlp_nodes): + model = Share(reference_name="MLP0") + MLP0 = MLP(nodes=mlp_nodes) + model._reference = MLP0 + + phlower_tensors = phlower_tensor_collection( + {"sample_input": PhlowerTensor(tensor=torch.rand(3, mlp_nodes[0]))} + ) + + mlp_val = MLP0(phlower_tensors) + model_val = model(phlower_tensors) + + np.testing.assert_array_almost_equal( + mlp_val.to_tensor().detach(), model_val.to_tensor().detach() + ) + + +@pytest.mark.parametrize("reference_name", ["MLP0", "GCP0"]) +def test__search_reference_name(reference_name): + model = Share(reference_name=reference_name) + + mocked = mock.MagicMock(IReadonlyReferenceGroup) + model.resolve(parent=mocked) + + mocked.search_module.assert_called_once_with(reference_name) diff --git a/tests/test_nn/test_group_module.py b/tests/test_nn/test_group_module.py new file mode 100644 index 0000000..21741a1 --- /dev/null +++ b/tests/test_nn/test_group_module.py @@ -0,0 +1,18 @@ +import pathlib + +import pytest + +from phlower.nn import PhlowerGroupModule +from phlower.settings import PhlowerSetting + +_SAMPLE_SETTING_DIR = pathlib.Path("tests/samples/settings") + + +@pytest.mark.parametrize("yaml_file", ["with_share_nn.yml"]) +def test__resolve_modules_from_setting(yaml_file): + setting_file = _SAMPLE_SETTING_DIR / yaml_file + + setting = PhlowerSetting.read_yaml(setting_file) + + setting.model.network.resolve(is_first=True) + _ = PhlowerGroupModule.from_setting(setting.model.network) diff --git a/tests/test_settings/test_module_settings/data/share_settings/check_gcn_share_nodes.yml b/tests/test_settings/test_module_settings/data/share_settings/check_gcn_share_nodes.yml new file mode 100644 index 0000000..a6a201b --- /dev/null +++ b/tests/test_settings/test_module_settings/data/share_settings/check_gcn_share_nodes.yml @@ -0,0 +1,50 @@ +misc: + + tests: + Share0: [5, 5] + +model: + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 5] + activations: ["Identity", "identity"] + + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + output_key: gcn0 + destinations: + - Share0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 + + - nn_type: Share + name: Share0 + input_keys: + - gcn0 + output_key: out_feature0 + nn_parameters: + reference_name: GCN0 diff --git a/tests/test_settings/test_module_settings/data/share_settings/check_mlp_share_nodes.yml b/tests/test_settings/test_module_settings/data/share_settings/check_mlp_share_nodes.yml new file mode 100644 index 0000000..064546f --- /dev/null +++ b/tests/test_settings/test_module_settings/data/share_settings/check_mlp_share_nodes.yml @@ -0,0 +1,50 @@ +misc: + + tests: + Share0: [10, 20, 10] + +model: + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - Share0 + nn_parameters: + nodes: [-1, 20, 10] + activations: ["Identity", "identity"] + + - nn_type: Share + name: Share0 + input_keys: + - mlp0 + output_key: share1 + destinations: + - GCN0 + nn_parameters: + reference_name: MLP0 + + - nn_type: GCN + name: GCN0 + input_keys: + - share1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/test_module_settings/data/share_settings/with_share_nn.yml b/tests/test_settings/test_module_settings/data/share_settings/with_share_nn.yml new file mode 100644 index 0000000..abf8855 --- /dev/null +++ b/tests/test_settings/test_module_settings/data/share_settings/with_share_nn.yml @@ -0,0 +1,46 @@ + +model: + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - Share0 + nn_parameters: + nodes: [-1, 20, 10] + activations: ["Identity", "identity"] + + - nn_type: Share + name: Share0 + input_keys: + - mlp0 + output_key: share1 + destinations: + - GCN0 + nn_parameters: + reference_name: MLP0 + + - nn_type: GCN + name: GCN0 + input_keys: + - share1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/test_module_settings/test_share_settings.py b/tests/test_settings/test_module_settings/test_share_settings.py new file mode 100644 index 0000000..9355a13 --- /dev/null +++ b/tests/test_settings/test_module_settings/test_share_settings.py @@ -0,0 +1,30 @@ +import pathlib + +import pytest +import yaml + +from phlower.settings import PhlowerModelSetting, PhlowerSetting + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/share_settings" + + +def test__can_resolve(): + setting = PhlowerSetting.read_yaml(_TEST_DATA_DIR / "with_share_nn.yml") + setting.model.network.resolve(is_first=True) + + +@pytest.mark.parametrize( + "yaml_file", ["check_gcn_share_nodes.yml", "check_mlp_share_nodes.yml"] +) +def test__nodes_after_resolve(yaml_file): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + + assert len(content["misc"]["tests"].items()) > 0 + + for key, value in content["misc"]["tests"].items(): + target = setting.network.search_module_setting(key) + assert target.get_n_nodes() == value From ce73286397977fd7d9b5a1e1fa6719163a961838 Mon Sep 17 00:00:00 2001 From: Riku Sakamoto Date: Wed, 7 Aug 2024 16:15:30 +0900 Subject: [PATCH 05/89] fix default parameter. (MLP, bias) --- src/phlower/nn/_core_modules/_gcn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index d12a546..57b5e15 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -41,7 +41,7 @@ def __init__( dropouts: list[float] | None = None, repeat: int = 1, factor: float = 1.0, - bias: bool = False, + bias: bool = True, ) -> None: super().__init__() From 1219392053b4830b1663c8814c3fc66dcbdd7673 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 7 Aug 2024 17:31:46 +0900 Subject: [PATCH 06/89] add tests for forward operation in groupModule. add reference info in mermaid state diagram. --- src/phlower/nn/_core_modules/_concatenator.py | 5 +- src/phlower/nn/_core_modules/_gcn.py | 5 +- src/phlower/nn/_core_modules/_mlp.py | 5 +- src/phlower/nn/_core_modules/_share.py | 5 +- src/phlower/nn/_group_module.py | 10 +++- src/phlower/nn/_interface_module.py | 10 +++- src/phlower/nn/_phlower_module_adpter.py | 14 +++++- tests/test_nn/out/.gitignore | 2 + tests/test_nn/test_group_module.py | 50 ++++++++++++++++++- 9 files changed, 95 insertions(+), 11 deletions(-) create mode 100644 tests/test_nn/out/.gitignore diff --git a/src/phlower/nn/_core_modules/_concatenator.py b/src/phlower/nn/_core_modules/_concatenator.py index 2dbbfae..4fc600b 100644 --- a/src/phlower/nn/_core_modules/_concatenator.py +++ b/src/phlower/nn/_core_modules/_concatenator.py @@ -38,7 +38,7 @@ def get_nn_name(cls) -> str: return "Concatenator" @classmethod - def need_resolve(cls) -> bool: + def need_reference(cls) -> bool: return False def __init__(self, activation: str, nodes: list[int] = None): @@ -51,6 +51,9 @@ def resolve( self, *, parent: IReadonlyReferenceGroup | None = None, **kwards ) -> None: ... + def get_reference_name(self) -> str | None: + return None + def forward( self, data: IPhlowerTensorCollections, diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index 6061e3d..ee542c9 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -37,7 +37,7 @@ def get_nn_name(cls) -> str: return "GCN" @classmethod - def need_resolve(cls) -> bool: + def need_reference(cls) -> bool: return False def __init__( @@ -70,6 +70,9 @@ def resolve( self, *, parent: IReadonlyReferenceGroup | None = None, **kwards ) -> None: ... + def get_reference_name(self) -> str | None: + return None + def forward( self, data: IPhlowerTensorCollections, diff --git a/src/phlower/nn/_core_modules/_mlp.py b/src/phlower/nn/_core_modules/_mlp.py index bd0074c..60f28e5 100644 --- a/src/phlower/nn/_core_modules/_mlp.py +++ b/src/phlower/nn/_core_modules/_mlp.py @@ -38,7 +38,7 @@ def get_nn_name(cls) -> str: return "MLP" @classmethod - def need_resolve(cls) -> bool: + def need_reference(cls) -> bool: return False def __init__( @@ -65,6 +65,9 @@ def resolve( self, *, parent: IReadonlyReferenceGroup | None = None, **kwards ) -> None: ... + def get_reference_name(self) -> str | None: + return None + def forward( self, data: IPhlowerTensorCollections, diff --git a/src/phlower/nn/_core_modules/_share.py b/src/phlower/nn/_core_modules/_share.py index 131ec4b..70fd7dd 100644 --- a/src/phlower/nn/_core_modules/_share.py +++ b/src/phlower/nn/_core_modules/_share.py @@ -27,7 +27,7 @@ def get_nn_name(cls) -> str: return "Share" @classmethod - def need_resolve(cls) -> bool: + def need_reference(cls) -> bool: return True def __init__(self, reference_name: str, **kwards) -> None: @@ -36,6 +36,9 @@ def __init__(self, reference_name: str, **kwards) -> None: self._reference_name = reference_name self._reference: IPhlowerCoreModule | None = None + def get_reference_name(self) -> str: + return self._reference_name + def resolve( self, *, parent: IReadonlyReferenceGroup | None = None, **kwards ) -> None: diff --git a/src/phlower/nn/_group_module.py b/src/phlower/nn/_group_module.py index a8907e8..4f74aad 100644 --- a/src/phlower/nn/_group_module.py +++ b/src/phlower/nn/_group_module.py @@ -14,6 +14,7 @@ ) from phlower.io._files import IPhlowerCheckpointFile from phlower.nn._interface_module import ( + IPhlowerCoreModule, IPhlowerModuleAdapter, IReadonlyReferenceGroup, ) @@ -159,10 +160,10 @@ def forward( def get_destinations(self) -> list[str]: return self._destinations - def search_module(self, name: str) -> IPhlowerModuleAdapter: + def search_module(self, name: str) -> IPhlowerCoreModule: for _module in self._phlower_modules: if _module.name == name: - return _module + return _module.get_core_module() raise KeyError(f"Module {name} is not found in group {self.name}.") @@ -183,3 +184,8 @@ def load_checkpoint_file( self.load_state_dict( content[TrainerSavedKeyType.MODEL_STATE_DICT.value] ) + + def get_core_module(self) -> IPhlowerCoreModule: + raise ValueError( + "`get_core_module` cannot be called in PhlowerGroupModule." + ) diff --git a/src/phlower/nn/_interface_module.py b/src/phlower/nn/_interface_module.py index e91a5fe..412aaac 100644 --- a/src/phlower/nn/_interface_module.py +++ b/src/phlower/nn/_interface_module.py @@ -21,7 +21,7 @@ def get_nn_name(cls) -> str: ... @classmethod @abc.abstractmethod - def need_resolve(cls) -> bool: ... + def need_reference(cls) -> bool: ... @abc.abstractmethod def resolve( @@ -36,10 +36,13 @@ def forward( supports: dict[str, PhlowerTensor] = None, ) -> PhlowerTensor: ... + @abc.abstractmethod + def get_reference_name(self) -> str | None: ... + class IReadonlyReferenceGroup(metaclass=abc.ABCMeta): @abc.abstractmethod - def search_module(self, name: str) -> IPhlowerModuleAdapter: ... + def search_module(self, name: str) -> IPhlowerCoreModule: ... class IPhlowerModuleAdapter(metaclass=abc.ABCMeta): @@ -72,3 +75,6 @@ def get_display_info(self) -> str: ... @abc.abstractmethod def draw(self, output_directory: pathlib.Path, recursive: bool): ... + + @abc.abstractmethod + def get_core_module(self) -> IPhlowerCoreModule: ... diff --git a/src/phlower/nn/_phlower_module_adpter.py b/src/phlower/nn/_phlower_module_adpter.py index 979b87c..90f0d87 100644 --- a/src/phlower/nn/_phlower_module_adpter.py +++ b/src/phlower/nn/_phlower_module_adpter.py @@ -61,7 +61,7 @@ def get_destinations(self) -> list[str]: def resolve( self, *, parent: IReadonlyReferenceGroup | None, **kwards ) -> None: - if not self._layer.need_resolve(): + if not self._layer.need_reference(): return self._layer.resolve(parent=parent, **kwards) @@ -70,9 +70,16 @@ def get_n_nodes(self) -> list[int]: return self._n_nodes def get_display_info(self) -> str: + if not self._layer.need_reference(): + return ( + f"nn_type: {self._layer.get_nn_name()}\n" + f"n_nodes: {self.get_n_nodes()}" + ) + return ( f"nn_type: {self._layer.get_nn_name()}\n" - f"n_nodes: {self.get_n_nodes()}" + f"n_nodes: {self.get_n_nodes()} \n" + f"reference: {self._layer.get_reference_name()}" ) def draw(self, output_directory: Path, recursive: bool): ... @@ -97,3 +104,6 @@ def forward( result = self._layer.forward(inputs, supports=supports) return phlower_tensor_collection({self._output_key: result}) + + def get_core_module(self) -> IPhlowerCoreModule: + return self._layer diff --git a/tests/test_nn/out/.gitignore b/tests/test_nn/out/.gitignore new file mode 100644 index 0000000..c96a04f --- /dev/null +++ b/tests/test_nn/out/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore \ No newline at end of file diff --git a/tests/test_nn/test_group_module.py b/tests/test_nn/test_group_module.py index 21741a1..5a39c92 100644 --- a/tests/test_nn/test_group_module.py +++ b/tests/test_nn/test_group_module.py @@ -1,7 +1,11 @@ import pathlib +import numpy as np import pytest +import scipy.sparse as sp +from phlower._base import phlower_array +from phlower.collections import phlower_tensor_collection from phlower.nn import PhlowerGroupModule from phlower.settings import PhlowerSetting @@ -11,8 +15,52 @@ @pytest.mark.parametrize("yaml_file", ["with_share_nn.yml"]) def test__resolve_modules_from_setting(yaml_file): setting_file = _SAMPLE_SETTING_DIR / yaml_file - setting = PhlowerSetting.read_yaml(setting_file) setting.model.network.resolve(is_first=True) _ = PhlowerGroupModule.from_setting(setting.model.network) + + +@pytest.mark.parametrize("yaml_file", ["with_share_nn.yml"]) +def test__draw(yaml_file): + output_directory = pathlib.Path(__file__).parent / "out" + + setting_file = _SAMPLE_SETTING_DIR / yaml_file + setting = PhlowerSetting.read_yaml(setting_file) + + setting.model.network.resolve(is_first=True) + group = PhlowerGroupModule.from_setting(setting.model.network) + + group.draw(output_directory) + + +@pytest.mark.parametrize( + "yaml_file, input_n_feature, n_nodes", [("with_share_nn.yml", 10, 20)] +) +def test__forward_and_backward(yaml_file, input_n_feature, n_nodes): + setting_file = _SAMPLE_SETTING_DIR / yaml_file + setting = PhlowerSetting.read_yaml(setting_file) + + setting.model.network.resolve(is_first=True) + group = PhlowerGroupModule.from_setting(setting.model.network) + + phlower_tensors = phlower_tensor_collection( + { + "feature0": phlower_array( + np.random.rand(n_nodes, input_n_feature).astype(np.float32) + ).to_phlower_tensor(), + "feature1": phlower_array( + np.random.rand(n_nodes, input_n_feature).astype(np.float32) + ).to_phlower_tensor(), + } + ) + + rng = np.random.default_rng() + sparse_adj = phlower_array( + sp.random( + n_nodes, n_nodes, density=0.1, random_state=rng, dtype=np.float32 + ) + ) + nodal_nadj = sparse_adj.to_phlower_tensor() + + _ = group.forward(data=phlower_tensors, supports={"support1": nodal_nadj}) From e67d9ae2d3ccaa8daba115712387aded45cf561a Mon Sep 17 00:00:00 2001 From: horiem Date: Thu, 8 Aug 2024 13:38:11 +0900 Subject: [PATCH 07/89] start adding en equivariant mlp --- .../nn/_core_modules/_en_equivariant_mlp.py | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 src/phlower/nn/_core_modules/_en_equivariant_mlp.py diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py new file mode 100644 index 0000000..42e320c --- /dev/null +++ b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import torch +from typing_extensions import Self + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._core_modules import _utils +from phlower.nn._interface_module import IPhlowerCoreModule +from phlower.settings._module_settings import MLPSetting + + +class EnEquivariantMLP(IPhlowerCoreModule, torch.nn.Module): + """E(n)-equivariant Multi Layer Perceptron""" + + @classmethod + def from_setting(cls, setting: MLPSetting) -> Self: + """Generate MLP from setting object + + Args: + setting (MLPSetting): setting object + + Returns: + Self: MLP object + """ + return EnEquivariantMLP(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + """Return neural network name + + Returns: + str: name + """ + return "EnEquivariantMLP" + + def __init__( + self, + nodes: list[int], + activations: list[str] | None = None, + dropouts: list[float] | None = None, + bias: bool = False, + ) -> None: + super().__init__() + + if activations is None: + activations = [] + if dropouts is None: + dropouts = [] + + self._chains = _utils.ExtendedLinearList( + nodes=nodes, activations=activations, dropouts=dropouts, bias=bias + ) + self._nodes = nodes + self._activations = activations + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + ) -> PhlowerTensor: + """forward function which overloads torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor], optional): + Graph object. Defaults to None. + + Returns: + PhlowerTensor: Tensor object + """ + h = data.unique_item() + for i in range(len(self._chains)): + h = self._chains.forward(h, index=i) + return h From 4071707d0e70562673dfe12424de1928853ba7b9 Mon Sep 17 00:00:00 2001 From: horiem Date: Thu, 8 Aug 2024 14:54:27 +0900 Subject: [PATCH 08/89] add flags and rank to PhlowerTensor --- .../_base/array/dense/_ndarray_wrapper.py | 4 +-- src/phlower/_base/tensors/_phlower_tensor.py | 30 ++++++++++++++++++- .../collections/arrays/_arrays_dict.py | 10 +++---- .../test_tensors/test__phlower_tensor.py | 24 +++++++++++++++ 4 files changed, 60 insertions(+), 8 deletions(-) diff --git a/src/phlower/_base/array/dense/_ndarray_wrapper.py b/src/phlower/_base/array/dense/_ndarray_wrapper.py index 98c17fa..8522fb5 100644 --- a/src/phlower/_base/array/dense/_ndarray_wrapper.py +++ b/src/phlower/_base/array/dense/_ndarray_wrapper.py @@ -10,9 +10,9 @@ class NdArrayWrapper(IPhlowerArray): - def __init__(self, data: np.ndarray, is_timeseries: bool = False): + def __init__(self, data: np.ndarray, is_time_series: bool = False): self.data = data - self._is_time_series = is_timeseries + self._is_time_series = is_time_series @property def shape(self) -> tuple[int]: diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index b31c48b..f2bec45 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -27,6 +27,8 @@ def phlower_tensor( | dict[str, float] | None ) = None, + is_time_series: bool = False, + is_voxel: bool = False, ): if isinstance(tensor, PhlowerTensor): if dimension is not None: @@ -38,7 +40,9 @@ def phlower_tensor( dimension_tensor = _resolve_dimension_arg(dimension) - return PhlowerTensor(tensor=tensor, dimension_tensor=dimension_tensor) + return PhlowerTensor( + tensor=tensor, dimension_tensor=dimension_tensor, + is_time_series=is_time_series, is_voxel=is_voxel) def _resolve_dimension_arg( @@ -77,10 +81,14 @@ def __init__( self, tensor: torch.Tensor, dimension_tensor: PhlowerDimensionTensor | None = None, + is_time_series: bool = False, + is_voxel: bool = False, ): assert isinstance(tensor, torch.Tensor) self._tensor = tensor self._dimension_tensor = dimension_tensor + self._is_time_series = is_time_series + self._is_voxel = is_voxel @property def has_dimension(self) -> bool: @@ -98,6 +106,14 @@ def shape(self) -> torch.Size: def is_sparse(self) -> bool: return self._tensor.layout == torch.sparse_coo + @property + def is_time_series(self) -> bool: + return self._is_time_series + + @property + def is_voxel(self) -> bool: + return self._is_voxel + def __str__(self) -> str: return ( f"PhysicsTensor({self._tensor}, " @@ -131,6 +147,18 @@ def coalesce(self) -> torch.Tensor: def size(self) -> torch.Size: return self._tensor.size() + def rank(self) -> int: + """Returns the tensor rank.""" + if self.is_sparse: + raise NotImplementedError + size = self.size() + start = 1 + if self.is_time_series: + start += 1 + if self.is_voxel: + start += 2 + return len(size[start:-1]) + def indices(self) -> torch.Tensor: return self._tensor.indices() diff --git a/src/phlower/collections/arrays/_arrays_dict.py b/src/phlower/collections/arrays/_arrays_dict.py index b2ca49e..1911e1b 100644 --- a/src/phlower/collections/arrays/_arrays_dict.py +++ b/src/phlower/collections/arrays/_arrays_dict.py @@ -20,7 +20,7 @@ def __init__(self, name: str, data: list[IPhlowerArray]) -> None: assert len(self._data) > 0 self._is_sparse = self._reduce_is_sparse() - self._is_timeseries = self._reduce_is_timeseries() + self._is_time_series = self._reduce_is_time_series() def _reduce_is_sparse(self): _is_sparse = np.unique(np.array([v.is_sparse for v in self._data])) @@ -31,7 +31,7 @@ def _reduce_is_sparse(self): ) return _is_sparse.item() - def _reduce_is_timeseries(self): + def _reduce_is_time_series(self): _is_time_series = np.unique( np.array([v.is_time_series for v in self._data]) ) @@ -46,8 +46,8 @@ def __len__(self) -> int: return len(self._data) @property - def is_timeseries(self) -> bool: - return self._is_timeseries + def is_time_series(self) -> bool: + return self._is_time_series @property def is_sparse(self) -> bool: @@ -69,7 +69,7 @@ def to_batched_tensor( if self.is_sparse: return to_batch(tensors) - if self.is_timeseries: + if self.is_time_series: return to_batch(tensors, dense_concat_dim=1) return to_batch(tensors, dense_concat_dim=0) diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index 5511817..e34e324 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -97,3 +97,27 @@ def test__getitem(key): np.testing.assert_array_almost_equal( phlower_tensor[key].to_tensor(), torch_tensor[key] ) + + +@pytest.mark.parametrize( + "is_time_series, is_voxel, size, desired_rank", + [ + (False, False, [100, 16], 0), + (False, False, [100, 3, 16], 1), + (False, False, [100, 3, 3, 16], 2), + ( True, False, [4, 100, 16], 0), + ( True, False, [4, 100, 3, 16], 1), + ( True, False, [4, 100, 3, 3, 16], 2), + (False, True, [10, 10, 10, 16], 0), + (False, True, [10, 10, 10, 3, 16], 1), + (False, True, [10, 10, 10, 3, 3, 16], 2), + ( True, True, [4, 10, 10, 10, 16], 0), + ( True, True, [4, 10, 10, 10, 3, 16], 1), + ( True, True, [4, 10, 10, 10, 3, 3, 16], 2), + ], +) +def test__rank(is_time_series, is_voxel, size, desired_rank): + torch_tensor = torch.rand(*size) + phlower_tensor = PhlowerTensor( + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + assert phlower_tensor.rank() == desired_rank From fa28ef6fc8984e2f8c0f14587bf98e9387c06c32 Mon Sep 17 00:00:00 2001 From: horiem Date: Thu, 8 Aug 2024 18:40:13 +0900 Subject: [PATCH 09/89] add optional args --- src/phlower/_base/array/_interface_wrapper.py | 2 ++ .../_base/array/dense/_ndarray_wrapper.py | 5 ++++- .../_base/array/sparse/_sparse_array_wrapper.py | 6 +++++- src/phlower/_base/tensors/_interface.py | 17 ++++++++++------- src/phlower/_base/tensors/_phlower_tensor.py | 17 ++++++----------- src/phlower/collections/arrays/_arrays_dict.py | 5 ++++- src/phlower/utils/exceptions.py | 4 ++++ .../test_tensors/test__phlower_tensor.py | 12 +++++++++++- 8 files changed, 46 insertions(+), 22 deletions(-) diff --git a/src/phlower/_base/array/_interface_wrapper.py b/src/phlower/_base/array/_interface_wrapper.py index dbea74a..2d71ceb 100644 --- a/src/phlower/_base/array/_interface_wrapper.py +++ b/src/phlower/_base/array/_interface_wrapper.py @@ -72,6 +72,8 @@ def to_phlower_tensor( device: str | torch.device | None = None, non_blocking: bool = False, dimension: PhysicalDimensions | None = None, + is_time_series: bool = False, + is_voxel: bool = False, ) -> PhlowerTensor: ... @abc.abstractmethod diff --git a/src/phlower/_base/array/dense/_ndarray_wrapper.py b/src/phlower/_base/array/dense/_ndarray_wrapper.py index 8522fb5..1d14ce6 100644 --- a/src/phlower/_base/array/dense/_ndarray_wrapper.py +++ b/src/phlower/_base/array/dense/_ndarray_wrapper.py @@ -64,9 +64,12 @@ def to_phlower_tensor( device: str | torch.device | None = None, non_blocking: bool = False, dimension: PhysicalDimensions | None = None, + is_time_series: bool = False, + is_voxel: bool = False, ) -> PhlowerTensor: _tensor = phlower_tensor( - tensor=torch.from_numpy(self.data), dimension=dimension + tensor=torch.from_numpy(self.data), dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel, ) _tensor.to(device=device, non_blocking=non_blocking) return _tensor diff --git a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py index 9205264..ea8c75b 100644 --- a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py +++ b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py @@ -96,6 +96,8 @@ def to_phlower_tensor( device: str | torch.device | None = None, non_blocking: bool = False, dimension: PhysicalDimensions | None = None, + is_time_series: bool = False, + is_voxel: bool = False, ) -> PhlowerTensor: sparse_tensor = torch.sparse_coo_tensor( torch.stack( @@ -107,7 +109,9 @@ def to_phlower_tensor( torch.from_numpy(self._sparse_data.data), self._sparse_data.shape, ) - _tensor = phlower_tensor(tensor=sparse_tensor, dimension=dimension) + _tensor = phlower_tensor( + tensor=sparse_tensor, dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) _tensor = _tensor.coalesce() _tensor.to(device=device, non_blocking=non_blocking) return _tensor diff --git a/src/phlower/_base/tensors/_interface.py b/src/phlower/_base/tensors/_interface.py index af736d3..fbe5f85 100644 --- a/src/phlower/_base/tensors/_interface.py +++ b/src/phlower/_base/tensors/_interface.py @@ -1,7 +1,7 @@ from __future__ import annotations import abc -from collections.abc import Callable, Sequence +from collections.abc import Callable import torch @@ -21,6 +21,12 @@ def shape(self) -> torch.Size: ... @abc.abstractproperty def is_sparse(self) -> bool: ... + @abc.abstractproperty + def is_time_series(self) -> bool: ... + + @abc.abstractproperty + def is_voxel(self) -> bool: ... + @abc.abstractmethod def values(self) -> torch.Tensor: ... @@ -46,16 +52,13 @@ def to_tensor(self) -> torch.Tensor: ... def size(self) -> torch.Size: ... @abc.abstractmethod - def indices(self) -> torch.Tensor: ... - - @abc.abstractmethod - def coalesce(self) -> IPhlowerTensor: ... + def rank(self) -> int: ... @abc.abstractmethod - def reshape(self, shape: Sequence[int]) -> IPhlowerTensor: ... + def indices(self) -> torch.Tensor: ... @abc.abstractmethod - def slice(self, slice_range: tuple[slice, ...]) -> IPhlowerTensor: ... + def coalesce(self) -> IPhlowerTensor: ... @abc.abstractmethod def to(self, device: str, non_blocking: bool = False) -> None: ... diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index f2bec45..177efef 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Iterable from typing import Any import torch @@ -13,7 +13,10 @@ ) from phlower._base.tensors._interface import IPhlowerTensor from phlower.utils import get_logger -from phlower.utils.exceptions import DimensionIncompatibleError +from phlower.utils.exceptions import ( + DimensionIncompatibleError, + PhlowerSparseRankUndefinedError, +) logger = get_logger(__name__) @@ -150,7 +153,7 @@ def size(self) -> torch.Size: def rank(self) -> int: """Returns the tensor rank.""" if self.is_sparse: - raise NotImplementedError + raise PhlowerSparseRankUndefinedError size = self.size() start = 1 if self.is_time_series: @@ -165,14 +168,6 @@ def indices(self) -> torch.Tensor: def values(self) -> torch.Tensor: return self._tensor.values() - def reshape(self, shape: Sequence[int]) -> PhlowerTensor: - self._tensor = self._tensor.reshape(shape) - return self - - def slice(self, slice_range: tuple[slice, ...]) -> PhlowerTensor: - tmp = self._tensor[slice_range] - return PhlowerTensor(tmp, dimension_tensor=self._dimension_tensor) - def to(self, device: str, non_blocking: bool = False) -> None: self._tensor.to(device, non_blocking=non_blocking) if self.has_dimension: diff --git a/src/phlower/collections/arrays/_arrays_dict.py b/src/phlower/collections/arrays/_arrays_dict.py index 1911e1b..6d735d5 100644 --- a/src/phlower/collections/arrays/_arrays_dict.py +++ b/src/phlower/collections/arrays/_arrays_dict.py @@ -58,10 +58,13 @@ def to_batched_tensor( device: str | torch.device | None = None, non_blocking: bool = False, dimensions: PhysicalDimensions | None = None, + is_time_series: bool = False, + is_voxel: bool = False, ): tensors = [ v.to_phlower_tensor( - device=device, non_blocking=non_blocking, dimension=dimensions + device=device, non_blocking=non_blocking, dimension=dimensions, + is_time_series=is_time_series, is_voxel=is_voxel, ) for v in self._data ] diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index fde6954..4b6e356 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -41,3 +41,7 @@ class PhlowerFeatureStoreOverwriteError(ValueError): """This error raises when to try to overwrite feature store item""" ... + + +class PhlowerSparseRankUndefinedError(ValueError): + """This error raises when trying to access tensor rank for sparse tensor""" diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index e34e324..d21c73d 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -3,7 +3,10 @@ import torch from phlower import PhlowerTensor, phlower_dimension_tensor, phlower_tensor -from phlower.utils.exceptions import DimensionIncompatibleError +from phlower.utils.exceptions import ( + DimensionIncompatibleError, + PhlowerSparseRankUndefinedError, +) def test__add(): @@ -121,3 +124,10 @@ def test__rank(is_time_series, is_voxel, size, desired_rank): phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) assert phlower_tensor.rank() == desired_rank + + +def test__raises_phlower_sparse_rank_undefined_error(): + torch_sparse_tensor = torch.eye(5).to_sparse() + phlower_sparse_tensor = PhlowerTensor(torch_sparse_tensor) + with pytest.raises(PhlowerSparseRankUndefinedError): + phlower_sparse_tensor.rank() From 217ee00ab7d9ddd3dda6cfe07e0273bd298eaed2 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 00:31:38 +0900 Subject: [PATCH 10/89] update tests --- .../test_tensors/test__phlower_tensor.py | 112 +++++++++++++++--- .../test_core_modules/test_functions.py | 47 ++++++++ tests/test_nn/test_core_modules/test_gcn.py | 41 +++++++ tests/test_settings/test_phlower_setting.py | 5 +- 4 files changed, 186 insertions(+), 19 deletions(-) create mode 100644 tests/test_nn/test_core_modules/test_functions.py create mode 100644 tests/test_nn/test_core_modules/test_gcn.py diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index d21c73d..2d54fe5 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -5,7 +5,7 @@ from phlower import PhlowerTensor, phlower_dimension_tensor, phlower_tensor from phlower.utils.exceptions import ( DimensionIncompatibleError, - PhlowerSparseRankUndefinedError, + PhlowerSparseUnsupportedError, ) @@ -86,22 +86,6 @@ def test__tanh(): np.testing.assert_array_almost_equal(cp.to_tensor(), c) -@pytest.mark.parametrize( - "key", - [ - 3, - [1, 3, 4], - [True, False, False, True, True], - ], -) -def test__getitem(key): - torch_tensor = torch.rand(5) - phlower_tensor = PhlowerTensor(torch_tensor) - np.testing.assert_array_almost_equal( - phlower_tensor[key].to_tensor(), torch_tensor[key] - ) - - @pytest.mark.parametrize( "is_time_series, is_voxel, size, desired_rank", [ @@ -126,8 +110,100 @@ def test__rank(is_time_series, is_voxel, size, desired_rank): assert phlower_tensor.rank() == desired_rank +@pytest.mark.parametrize( + "is_time_series, is_voxel, size, desired_n_vertices", + [ + (False, False, [100, 16], 100), + (False, False, [100, 3, 16], 100), + (False, False, [100, 3, 3, 16], 100), + ( True, False, [4, 100, 16], 100), + ( True, False, [4, 100, 3, 16], 100), + ( True, False, [4, 100, 3, 3, 16], 100), + (False, True, [10, 10, 10, 16], 1000), + (False, True, [10, 10, 10, 3, 16], 1000), + (False, True, [10, 10, 10, 3, 3, 16], 1000), + ( True, True, [4, 10, 10, 10, 16], 1000), + ( True, True, [4, 10, 10, 10, 3, 16], 1000), + ( True, True, [4, 10, 10, 10, 3, 3, 16], 1000), + ], +) +def test__n_vertices(is_time_series, is_voxel, size, desired_n_vertices): + torch_tensor = torch.rand(*size) + phlower_tensor = PhlowerTensor( + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + assert phlower_tensor.n_vertices() == desired_n_vertices + + def test__raises_phlower_sparse_rank_undefined_error(): torch_sparse_tensor = torch.eye(5).to_sparse() phlower_sparse_tensor = PhlowerTensor(torch_sparse_tensor) - with pytest.raises(PhlowerSparseRankUndefinedError): + with pytest.raises(PhlowerSparseUnsupportedError): phlower_sparse_tensor.rank() + + +@pytest.mark.parametrize( + "is_time_series, is_voxel, size, desired_shape", + [ + (False, False, [100, 16], (100, 16)), + (False, False, [100, 3, 16], (100, 3 * 16)), + (False, False, [100, 3, 3, 16], (100, 3 * 3 * 16)), + ( True, False, [4, 100, 16], (100, 4 * 16)), + ( True, False, [4, 100, 3, 16], (100, 4 * 3 * 16)), + ( True, False, [4, 100, 3, 3, 16], (100, 4 * 3 * 3 * 16)), + (False, True, [10, 10, 10, 16], (1000, 16)), + (False, True, [10, 10, 10, 3, 16], (1000, 3 * 16)), + (False, True, [10, 10, 10, 3, 3, 16], (1000, 3 * 3 * 16)), + ( True, True, [4, 10, 10, 10, 16], (1000, 4 * 16)), + ( True, True, [4, 10, 10, 10, 3, 16], (1000, 4 * 3 * 16)), + ( True, True, [4, 10, 10, 10, 3, 3, 16], (1000, 4 * 3 * 3 * 16)), + ], +) +def test__to_2d(is_time_series, is_voxel, size, desired_shape): + torch_tensor = torch.rand(*size) + phlower_tensor = PhlowerTensor( + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + assert phlower_tensor.to_2d().shape == desired_shape + + +@pytest.mark.parametrize( + "is_time_series, is_voxel, size", + [ + (False, False, [100, 16]), + (False, False, [100, 3, 16]), + (False, False, [100, 3, 3, 16]), + ( True, False, [4, 100, 16]), + ( True, False, [4, 100, 3, 16]), + ( True, False, [4, 100, 3, 3, 16]), + (False, True, [10, 10, 10, 16]), + (False, True, [10, 10, 10, 3, 16]), + (False, True, [10, 10, 10, 3, 3, 16]), + ( True, True, [4, 10, 10, 10, 16]), + ( True, True, [4, 10, 10, 10, 3, 16]), + ( True, True, [4, 10, 10, 10, 3, 3, 16]), + ], +) +def test__to_from_2d(is_time_series, is_voxel, size): + torch_tensor = torch.rand(*size) + phlower_tensor = PhlowerTensor( + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + actual = phlower_tensor.to_2d().from_2d() + np.testing.assert_almost_equal( + actual.to_tensor().numpy(), + phlower_tensor.to_tensor().numpy()) + + +@pytest.mark.parametrize( + "input_shape, pattern, dict_shape, desired_shape", + [ + ((10, 3, 16), "n p a -> n (p a)", {"a": 16}, (10, 3 * 16)), + ((10, 3 * 16), "n (p a) -> n p a", {"p": 3}, (10, 3, 16)), + ], +) +def test__rearrange(input_shape, pattern, dict_shape, desired_shape): + phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) + actual = phlower_tensor.rearrange(pattern, **dict_shape) + assert actual.shape == desired_shape + assert actual.dict_shape == dict_shape + + actual_inversed = actual.inverse_rearrange() + assert actual_inversed.shape == phlower_tensor.shape diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py new file mode 100644 index 0000000..881d481 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -0,0 +1,47 @@ + +import numpy as np +import pytest +import torch +from scipy import sparse as sp + +from phlower import PhlowerTensor +from phlower.nn._core_modules import _functions + + +@pytest.mark.parametrize( + "size, is_time_series", + [ + ((10, 1), False), + ((10, 16), False), + ((10, 3, 16), False), + ((4, 10, 1), True), + ((4, 10, 16), True), + ((4, 10, 3, 16), True), + ], +) +def test__spmm(size, is_time_series): + phlower_tensor = PhlowerTensor( + torch.rand(*size), is_time_series=is_time_series) + n = phlower_tensor.n_vertices() + sparse = PhlowerTensor(torch.rand(n, n).to_sparse()) + + actual_spmm = _functions.spmm(sparse, phlower_tensor).to_tensor().numpy() + sp_sparse = sp.coo_array(sparse.to_tensor().to_dense().numpy()) + np_dense = phlower_tensor.to_tensor().numpy() + + def assert_correct(actual, array): + dim_feat = len(array.shape) - 1 + if dim_feat == 1: + desired = sp_sparse @ array + np.testing.assert_almost_equal(actual, desired) + return + + for i in range(array.shape[1]): + assert_correct(actual[:, i], array[:, i]) + return + + if is_time_series: + for t in range(size[0]): + assert_correct(actual_spmm[t], np_dense[t]) + else: + assert_correct(actual_spmm, np_dense) diff --git a/tests/test_nn/test_core_modules/test_gcn.py b/tests/test_nn/test_core_modules/test_gcn.py new file mode 100644 index 0000000..29b9b73 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_gcn.py @@ -0,0 +1,41 @@ +import pytest +import torch + +from phlower import PhlowerTensor +from phlower.collections import phlower_tensor_collection +from phlower.nn import GCN + + +def test__can_call_parameters(): + model = GCN(nodes=[4, 8], support_name='support') + + # To check Concatenator inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "size, is_time_series", + [ + ((10, 1), False), + ((10, 16), False), + ((10, 3, 16), False), + ((4, 10, 1), True), + ((4, 10, 16), True), + ((4, 10, 3, 16), True), + ], +) +def test__concatenated_tensor_shape(size, is_time_series): + phlower_tensor = PhlowerTensor( + torch.rand(*size), is_time_series=is_time_series) + phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) + n = phlower_tensor.n_vertices() + dict_supports = {'support': PhlowerTensor(torch.rand(n, n).to_sparse())} + + model = GCN( + nodes=[size[-1], size[-1]], + support_name='support', activations=['tanh']) + + actual = model(phlower_tensors, supports=dict_supports) + + assert actual.shape == size + assert actual.is_time_series == phlower_tensor.is_time_series diff --git a/tests/test_settings/test_phlower_setting.py b/tests/test_settings/test_phlower_setting.py index f262f15..350424f 100644 --- a/tests/test_settings/test_phlower_setting.py +++ b/tests/test_settings/test_phlower_setting.py @@ -1,4 +1,5 @@ import pathlib +import shutil import pydantic import pytest @@ -46,8 +47,10 @@ def test__model_dump(): "tests/test_settings/data/e2e/setting1.yml" ) + output_directory = pathlib.Path("tests/test_settings/tmp") + shutil.rmtree(output_directory, ignore_errors=True) _ = PhlowerYamlFile.save( - output_directory=pathlib.Path("tests/test_settings/tmp"), + output_directory=output_directory, file_basename="output", data=setting.model_dump(), ) From 34dc21d3bcff1b7ef63a9e01cc82fc7dc3b7dee6 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 00:31:45 +0900 Subject: [PATCH 11/89] add einops --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 898b77f..324f2ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ pipe = "^2.2" pyvista = "^0.43.10" tqdm = "^4.66.4" pandas = "^2.2.2" +einops = "^0.8.0" [tool.poetry.group.dev.dependencies] pytest = "^8.0.2" From d3bec2fa4430f7f110960cf1b5c6fcdbff3f2230 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 00:33:14 +0900 Subject: [PATCH 12/89] add reshape and rearrange --- src/phlower/_base/tensors/_interface.py | 6 +- src/phlower/_base/tensors/_phlower_tensor.py | 140 +++++++++++++++++- .../tensors/_unsupported_function_names.py | 5 + 3 files changed, 141 insertions(+), 10 deletions(-) create mode 100644 src/phlower/_base/tensors/_unsupported_function_names.py diff --git a/src/phlower/_base/tensors/_interface.py b/src/phlower/_base/tensors/_interface.py index fbe5f85..28e8909 100644 --- a/src/phlower/_base/tensors/_interface.py +++ b/src/phlower/_base/tensors/_interface.py @@ -42,9 +42,6 @@ def __mul__(self, other) -> IPhlowerTensor: ... @abc.abstractmethod def __len__(self) -> int: ... - @abc.abstractmethod - def __getitem__(self, key) -> IPhlowerTensor: ... - @abc.abstractmethod def to_tensor(self) -> torch.Tensor: ... @@ -54,6 +51,9 @@ def size(self) -> torch.Size: ... @abc.abstractmethod def rank(self) -> int: ... + @abc.abstractmethod + def n_vertices(self) -> int: ... + @abc.abstractmethod def indices(self) -> torch.Tensor: ... diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 177efef..6e1cd93 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -1,8 +1,10 @@ from __future__ import annotations -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence from typing import Any +import einops +import numpy as np import torch from pipe import select @@ -12,10 +14,15 @@ phlower_dimension_tensor, ) from phlower._base.tensors._interface import IPhlowerTensor +from phlower._base.tensors._unsupported_function_names import ( + UNSUPPORTED_FUNCTION_NAMES, +) from phlower.utils import get_logger from phlower.utils.exceptions import ( DimensionIncompatibleError, - PhlowerSparseRankUndefinedError, + PhlowerReshapeError, + PhlowerSparseUnsupportedError, + PhlowerUnsupportedTorchFunctionError, ) logger = get_logger(__name__) @@ -86,12 +93,18 @@ def __init__( dimension_tensor: PhlowerDimensionTensor | None = None, is_time_series: bool = False, is_voxel: bool = False, + original_pattern: str | None = None, + current_pattern: str | None = None, + dict_shape: dict[str, int] | None = None, ): assert isinstance(tensor, torch.Tensor) self._tensor = tensor self._dimension_tensor = dimension_tensor self._is_time_series = is_time_series self._is_voxel = is_voxel + self._original_pattern = original_pattern + self._current_pattern = current_pattern + self._dict_shape = dict_shape @property def has_dimension(self) -> bool: @@ -117,9 +130,21 @@ def is_time_series(self) -> bool: def is_voxel(self) -> bool: return self._is_voxel + @property + def original_pattern(self) -> str | None: + return self._original_pattern + + @property + def current_pattern(self) -> str | None: + return self._current_pattern + + @property + def dict_shape(self) -> dict[str, int] | None: + return self._dict_shape + def __str__(self) -> str: return ( - f"PhysicsTensor({self._tensor}, " + f"PhlowerTensor({self._tensor}, " f"Dimension: {self._dimension_tensor})" ) @@ -138,9 +163,6 @@ def __mul__(self, other) -> PhlowerTensor: def __len__(self) -> int: return len(self._tensor) - def __getitem__(self, key: Any) -> PhlowerTensor: - return PhlowerTensor(self._tensor[key], self._dimension_tensor) - def to_tensor(self) -> torch.Tensor: return self._tensor @@ -153,7 +175,8 @@ def size(self) -> torch.Size: def rank(self) -> int: """Returns the tensor rank.""" if self.is_sparse: - raise PhlowerSparseRankUndefinedError + raise PhlowerSparseUnsupportedError( + "Cannot call rank() for sparse PhlowerTensor") size = self.size() start = 1 if self.is_time_series: @@ -162,12 +185,112 @@ def rank(self) -> int: start += 2 return len(size[start:-1]) + def n_vertices(self) -> int: + """Returns the number of vertices.""" + if self.is_sparse: + raise PhlowerSparseUnsupportedError( + "Cannot call n_vertices() for sparse PhlowerTensor") + size = self.size() + start = 0 + if self.is_time_series: + start += 1 + + if self.is_voxel: + return np.prod(size[start:start+3]) + + return size[start] + def indices(self) -> torch.Tensor: return self._tensor.indices() def values(self) -> torch.Tensor: return self._tensor.values() + def to_2d(self) -> PhlowerTensor: + """Convert to 2D tensor which has (n_vertices, -1) shape""" + shape = self.shape + dict_shape = {} + + space_start = 0 + if self.is_time_series: + t_pattern = "t " + space_start += 1 + dict_shape.update({"t": shape[0]}) + else: + t_pattern = "" + if self.is_voxel: + space_pattern = "x y z " + dict_shape.update({ + "x": shape[space_start], + "y": shape[space_start + 1], + "z": shape[space_start + 2], + }) + feat_start = space_start + 3 + else: + space_pattern = "n " + dict_shape.update({"n": shape[space_start]}) + feat_start = space_start + 1 + + feat_pattern = " ".join([ + f"a{i}" for i in range(len(shape[feat_start:]))]) + # Do not include the last axis in case modified by NNs. + dict_shape.update({ + f"a{i}": s for i, s in enumerate(shape[feat_start:-1])}) + + original_pattern = f"{t_pattern}{space_pattern}{feat_pattern}" + current_pattern = f"({space_pattern}) ({t_pattern}{feat_pattern})" + tensor_2d = einops.rearrange( + self.to_tensor(), f"{original_pattern} -> {current_pattern}") + return PhlowerTensor( + tensor_2d, dimension_tensor=self.dimension, + is_time_series=False, is_voxel=False, + original_pattern=original_pattern, + current_pattern=current_pattern, + dict_shape=dict_shape) + + def from_2d(self) -> PhlowerTensor: + if self.original_pattern is None: + raise PhlowerReshapeError( + "No original_pattern found. Run to_2d first.") + is_time_series = "t" in self.original_pattern + is_voxel = "x" in self.original_pattern + return self.inverse_rearrange( + is_time_series=is_time_series, is_voxel=is_voxel) + + def rearrange( + self, pattern: str, + is_time_series: bool = False, is_voxel: bool = False, + **kwargs) -> PhlowerTensor: + tensor = self.to_tensor() + original_pattern, current_pattern = pattern.split("->") + rearranged = einops.rearrange(tensor, pattern, **kwargs) + return PhlowerTensor( + rearranged, dimension_tensor=self.dimension, + is_time_series=is_time_series, is_voxel=is_voxel, + original_pattern=original_pattern, current_pattern=current_pattern, + dict_shape=kwargs) + + def inverse_rearrange( + self, is_time_series: bool = False, is_voxel: bool = False, + **kwargs) -> PhlowerTensor: + if self.original_pattern is None: + raise PhlowerReshapeError( + "No original_pattern found. Run rearrange first.") + pattern = f"{self.current_pattern} -> {self.original_pattern}" + kwargs.update(self.dict_shape) + return self.rearrange( + pattern, is_time_series=is_time_series, is_voxel=is_voxel, + **kwargs) + + def reshape( + self, shape: Sequence[int], + is_time_series: bool = False, is_voxel: bool = False, + ) -> PhlowerTensor: + return PhlowerTensor( + torch.reshape(self.to_tensor(), shape), + dimension_tensor=self.dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + def to(self, device: str, non_blocking: bool = False) -> None: self._tensor.to(device, non_blocking=non_blocking) if self.has_dimension: @@ -190,6 +313,9 @@ def __torch_function__( args: tuple, kwargs: dict | None = None, ): + if func.__name__ in UNSUPPORTED_FUNCTION_NAMES: + raise PhlowerUnsupportedTorchFunctionError( + f"Unsupported function: {func.__name__}") if kwargs is None: kwargs = {} diff --git a/src/phlower/_base/tensors/_unsupported_function_names.py b/src/phlower/_base/tensors/_unsupported_function_names.py new file mode 100644 index 0000000..c2ef632 --- /dev/null +++ b/src/phlower/_base/tensors/_unsupported_function_names.py @@ -0,0 +1,5 @@ + +UNSUPPORTED_FUNCTION_NAMES = [ + 'einsum', + 'reshape', +] From 20e02e9436ab1a71cf4b247620453028d78e3e45 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 00:33:27 +0900 Subject: [PATCH 13/89] add exceptions --- src/phlower/utils/exceptions.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 4b6e356..79ecb58 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -43,5 +43,22 @@ class PhlowerFeatureStoreOverwriteError(ValueError): ... -class PhlowerSparseRankUndefinedError(ValueError): - """This error raises when trying to access tensor rank for sparse tensor""" +class PhlowerSparseUnsupportedError(ValueError): + """ + This error raises when trying to call methods not supported for sparse + tensors + """ + + +class PhlowerUnsupportedTorchFunctionError(ValueError): + """ + This error raises when trying to call a function not supported + by the phlower library although torch does + """ + + +class PhlowerReshapeError(ValueError): + """ + This error raises when trying to reshape a tensor in an invalid + manner + """ From 46eb9323c1033b6d250df46f0deb068ce1848811 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 00:34:18 +0900 Subject: [PATCH 14/89] update gcn --- src/phlower/nn/_core_modules/_functions.py | 13 +++++++++++++ src/phlower/nn/_core_modules/_gcn.py | 17 ++++++++++++++--- 2 files changed, 27 insertions(+), 3 deletions(-) create mode 100644 src/phlower/nn/_core_modules/_functions.py diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py new file mode 100644 index 0000000..99ad326 --- /dev/null +++ b/src/phlower/nn/_core_modules/_functions.py @@ -0,0 +1,13 @@ + +import torch + +from phlower._base.tensors._interface import IPhlowerTensor + + +def spmm(sparse: IPhlowerTensor, x: IPhlowerTensor) -> IPhlowerTensor: + h = x.to_2d() + ret = torch.sparse.mm(sparse, h) + pattern = f"{h.current_pattern} -> {h.original_pattern}" + return ret.rearrange( + pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, + **h.dict_shape) diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index d12a546..c7310f3 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -87,8 +87,19 @@ def forward( def _propagate( self, x: PhlowerTensor, support: PhlowerTensor ) -> PhlowerTensor: - n_node = x.shape[0] - h = torch.reshape(x, (n_node, -1)) + # NOTE: Could be simplified as follows, + # but maybe slow due to reshape running every time + # from phlower.nn._core_modules import _utils + # h = x + # for _ in range(self._repeat): + # h = _functions.spmm(support, h) * self._factor + # return h + + h = x.to_2d() + pattern = f"{h.current_pattern} -> {h.original_pattern}" + dict_shape = h.dict_shape for _ in range(self._repeat): h = torch.sparse.mm(support, h) * self._factor - return torch.reshape(h, x.shape) + return h.rearrange( + pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, + **dict_shape) From dc17ec5fff028d1374129e91a83b26f4d97787a4 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 01:15:17 +0900 Subject: [PATCH 15/89] debug tensor to preserve flags --- src/phlower/_base/tensors/_phlower_tensor.py | 12 ++++++++++-- src/phlower/nn/_core_modules/_utils.py | 2 +- tests/test_nn/test_core_modules/test_gcn.py | 8 ++++---- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 6e1cd93..7ccb6f5 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -300,6 +300,7 @@ def detach(self) -> PhlowerTensor: return PhlowerTensor( self._tensor.detach(), dimension_tensor=self._dimension_tensor.detach(), + is_time_series=self.is_time_series, is_voxel=self.is_voxel, ) def backward(self) -> None: @@ -323,16 +324,23 @@ def __torch_function__( ret: torch.Tensor = func(*_tensors, **kwargs) + # NOTE: Assume flags for the first tensor is preserved + is_time_series = args[0].is_time_series + is_voxel = args[0].is_voxel + if not _has_dimension(args): # Unit calculation is not considered when unit tensor is not found. - return PhlowerTensor(ret) + return PhlowerTensor( + ret, is_time_series=is_time_series, is_voxel=is_voxel) _dimensions = _recursive_resolve( args, "_dimension_tensor", allow_none=False ) result_units = func(*_dimensions, **kwargs) - return PhlowerTensor(ret, result_units) + return PhlowerTensor( + ret, result_units, is_time_series=is_time_series, + is_voxel=is_voxel) def _recursive_resolve( diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index a5d8d3c..55a16e8 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -42,7 +42,7 @@ def __len__(self) -> int: def forward(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: assert index < self._n_chains - h = self._linears[index].forward(x) + h = self._linears[index](x) h = torch.nn.functional.dropout( h, p=self._dropouts[index], training=self.training ) diff --git a/tests/test_nn/test_core_modules/test_gcn.py b/tests/test_nn/test_core_modules/test_gcn.py index 29b9b73..53604a8 100644 --- a/tests/test_nn/test_core_modules/test_gcn.py +++ b/tests/test_nn/test_core_modules/test_gcn.py @@ -24,7 +24,7 @@ def test__can_call_parameters(): ((4, 10, 3, 16), True), ], ) -def test__concatenated_tensor_shape(size, is_time_series): +def test__gcn(size, is_time_series): phlower_tensor = PhlowerTensor( torch.rand(*size), is_time_series=is_time_series) phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) @@ -32,10 +32,10 @@ def test__concatenated_tensor_shape(size, is_time_series): dict_supports = {'support': PhlowerTensor(torch.rand(n, n).to_sparse())} model = GCN( - nodes=[size[-1], size[-1]], - support_name='support', activations=['tanh']) + nodes=[size[-1], size[-1], size[-1]], + support_name='support', activations=['tanh', 'identity']) actual = model(phlower_tensors, supports=dict_supports) assert actual.shape == size - assert actual.is_time_series == phlower_tensor.is_time_series + assert actual.is_time_series == is_time_series From 9dd6cda9c3fc608e14892c199d62ce34cb6b127c Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 01:43:13 +0900 Subject: [PATCH 16/89] debug tensor --- src/phlower/_base/tensors/_phlower_tensor.py | 36 ++++++++++++++------ 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 7ccb6f5..e0c662f 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -220,11 +220,14 @@ def to_2d(self) -> PhlowerTensor: t_pattern = "" if self.is_voxel: space_pattern = "x y z " - dict_shape.update({ - "x": shape[space_start], - "y": shape[space_start + 1], - "z": shape[space_start + 2], - }) + try: + dict_shape.update({ + "x": shape[space_start], + "y": shape[space_start + 1], + "z": shape[space_start + 2], + }) + except: + raise ValueError(self.shape, self.is_time_series, self.is_voxel) feat_start = space_start + 3 else: space_pattern = "n " @@ -325,8 +328,10 @@ def __torch_function__( ret: torch.Tensor = func(*_tensors, **kwargs) # NOTE: Assume flags for the first tensor is preserved - is_time_series = args[0].is_time_series - is_voxel = args[0].is_voxel + is_time_series = _recursive_resolve( + args, "_is_time_series", return_first_only=True) + is_voxel = _recursive_resolve( + args, "_is_voxel", return_first_only=True) if not _has_dimension(args): # Unit calculation is not considered when unit tensor is not found. @@ -344,12 +349,21 @@ def __torch_function__( def _recursive_resolve( - args: Iterable | Any, attr: str, allow_none: bool = True + args: Iterable | Any, attr: str, allow_none: bool = True, + return_first_only: bool = False, ) -> list[str]: if isinstance(args, tuple | list): - return [ - _recursive_resolve(v, attr, allow_none=allow_none) for v in args - ] + if return_first_only: + return _recursive_resolve( + args[0], attr, allow_none=allow_none, + return_first_only=return_first_only) + else: + return [ + _recursive_resolve( + v, attr, allow_none=allow_none, + return_first_only=return_first_only) + for v in args + ] _val = getattr(args, attr, args) From c9a216aae422a9fced805ec4047c4778829e66b0 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 9 Aug 2024 01:52:27 +0900 Subject: [PATCH 17/89] fix test tolerance --- tests/test_nn/test_core_modules/test_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 881d481..81ae3a9 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -33,7 +33,7 @@ def assert_correct(actual, array): dim_feat = len(array.shape) - 1 if dim_feat == 1: desired = sp_sparse @ array - np.testing.assert_almost_equal(actual, desired) + np.testing.assert_almost_equal(actual, desired, decimal=5) return for i in range(array.shape[1]): From 476e64b2e8cc5d00d85eac5db3bcad4da5880b4b Mon Sep 17 00:00:00 2001 From: horiem Date: Sat, 10 Aug 2024 17:18:52 +0900 Subject: [PATCH 18/89] remove unnecessary file --- .../nn/_core_modules/_en_equivariant_mlp.py | 77 ------------------- 1 file changed, 77 deletions(-) delete mode 100644 src/phlower/nn/_core_modules/_en_equivariant_mlp.py diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py deleted file mode 100644 index 42e320c..0000000 --- a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -import torch -from typing_extensions import Self - -from phlower._base.tensors import PhlowerTensor -from phlower.collections.tensors import IPhlowerTensorCollections -from phlower.nn._core_modules import _utils -from phlower.nn._interface_module import IPhlowerCoreModule -from phlower.settings._module_settings import MLPSetting - - -class EnEquivariantMLP(IPhlowerCoreModule, torch.nn.Module): - """E(n)-equivariant Multi Layer Perceptron""" - - @classmethod - def from_setting(cls, setting: MLPSetting) -> Self: - """Generate MLP from setting object - - Args: - setting (MLPSetting): setting object - - Returns: - Self: MLP object - """ - return EnEquivariantMLP(**setting.__dict__) - - @classmethod - def get_nn_name(cls) -> str: - """Return neural network name - - Returns: - str: name - """ - return "EnEquivariantMLP" - - def __init__( - self, - nodes: list[int], - activations: list[str] | None = None, - dropouts: list[float] | None = None, - bias: bool = False, - ) -> None: - super().__init__() - - if activations is None: - activations = [] - if dropouts is None: - dropouts = [] - - self._chains = _utils.ExtendedLinearList( - nodes=nodes, activations=activations, dropouts=dropouts, bias=bias - ) - self._nodes = nodes - self._activations = activations - - def forward( - self, - data: IPhlowerTensorCollections, - *, - supports: dict[str, PhlowerTensor] | None = None, - ) -> PhlowerTensor: - """forward function which overloads torch.nn.Module - - Args: - data (IPhlowerTensorCollections): - data which receives from predecessors - supports (dict[str, PhlowerTensor], optional): - Graph object. Defaults to None. - - Returns: - PhlowerTensor: Tensor object - """ - h = data.unique_item() - for i in range(len(self._chains)): - h = self._chains.forward(h, index=i) - return h From 32f766064436b4292293df663299faf4944fd1d8 Mon Sep 17 00:00:00 2001 From: horiem Date: Sat, 10 Aug 2024 17:19:05 +0900 Subject: [PATCH 19/89] update abstract method --- src/phlower/_base/tensors/_interface.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/phlower/_base/tensors/_interface.py b/src/phlower/_base/tensors/_interface.py index 28e8909..8476273 100644 --- a/src/phlower/_base/tensors/_interface.py +++ b/src/phlower/_base/tensors/_interface.py @@ -9,22 +9,28 @@ class IPhlowerTensor(metaclass=abc.ABCMeta): - @abc.abstractproperty + @property + @abc.abstractmethod def has_dimension(self) -> bool: ... - @abc.abstractproperty + @property + @abc.abstractmethod def dimension(self) -> PhlowerDimensionTensor | None: ... - @abc.abstractproperty + @property + @abc.abstractmethod def shape(self) -> torch.Size: ... - @abc.abstractproperty + @property + @abc.abstractmethod def is_sparse(self) -> bool: ... - @abc.abstractproperty + @property + @abc.abstractmethod def is_time_series(self) -> bool: ... - @abc.abstractproperty + @property + @abc.abstractmethod def is_voxel(self) -> bool: ... @abc.abstractmethod From 2b98b1945169a2701020274d4c8438c39b2c406c Mon Sep 17 00:00:00 2001 From: horiem Date: Sat, 10 Aug 2024 20:54:45 +0900 Subject: [PATCH 20/89] reflect time series for ndarray wrapper --- src/phlower/_base/array/dense/_ndarray_wrapper.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/phlower/_base/array/dense/_ndarray_wrapper.py b/src/phlower/_base/array/dense/_ndarray_wrapper.py index 1d14ce6..c7fe74d 100644 --- a/src/phlower/_base/array/dense/_ndarray_wrapper.py +++ b/src/phlower/_base/array/dense/_ndarray_wrapper.py @@ -64,9 +64,11 @@ def to_phlower_tensor( device: str | torch.device | None = None, non_blocking: bool = False, dimension: PhysicalDimensions | None = None, - is_time_series: bool = False, + is_time_series: bool | None = None, is_voxel: bool = False, ) -> PhlowerTensor: + if is_time_series is None: + is_time_series = self.is_time_series _tensor = phlower_tensor( tensor=torch.from_numpy(self.data), dimension=dimension, is_time_series=is_time_series, is_voxel=is_voxel, From b764310aa9685a91fcab50902060b0f212ec5e8c Mon Sep 17 00:00:00 2001 From: horiem Date: Sat, 10 Aug 2024 21:03:14 +0900 Subject: [PATCH 21/89] simplify tensor rearrange --- src/phlower/_base/tensors/_phlower_tensor.py | 84 ++++++------------- src/phlower/nn/_core_modules/_functions.py | 29 +++++-- src/phlower/nn/_core_modules/_gcn.py | 20 +---- .../test_tensors/test__phlower_tensor.py | 18 ++-- .../test_core_modules/test_functions.py | 29 ++++--- 5 files changed, 79 insertions(+), 101 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index e0c662f..73a2b6a 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -20,7 +20,6 @@ from phlower.utils import get_logger from phlower.utils.exceptions import ( DimensionIncompatibleError, - PhlowerReshapeError, PhlowerSparseUnsupportedError, PhlowerUnsupportedTorchFunctionError, ) @@ -93,18 +92,12 @@ def __init__( dimension_tensor: PhlowerDimensionTensor | None = None, is_time_series: bool = False, is_voxel: bool = False, - original_pattern: str | None = None, - current_pattern: str | None = None, - dict_shape: dict[str, int] | None = None, ): assert isinstance(tensor, torch.Tensor) self._tensor = tensor self._dimension_tensor = dimension_tensor self._is_time_series = is_time_series self._is_voxel = is_voxel - self._original_pattern = original_pattern - self._current_pattern = current_pattern - self._dict_shape = dict_shape @property def has_dimension(self) -> bool: @@ -130,18 +123,6 @@ def is_time_series(self) -> bool: def is_voxel(self) -> bool: return self._is_voxel - @property - def original_pattern(self) -> str | None: - return self._original_pattern - - @property - def current_pattern(self) -> str | None: - return self._current_pattern - - @property - def dict_shape(self) -> dict[str, int] | None: - return self._dict_shape - def __str__(self) -> str: return ( f"PhlowerTensor({self._tensor}, " @@ -206,8 +187,20 @@ def indices(self) -> torch.Tensor: def values(self) -> torch.Tensor: return self._tensor.values() - def to_2d(self) -> PhlowerTensor: - """Convert to 2D tensor which has (n_vertices, -1) shape""" + def to_vertexwise(self) -> PhlowerTensor: + """ + Convert to vertexwise 2D tensor which has (n_vertices, -1) shape. + + Returns: + vertexwise_tensor : PhlowerTensor + Vertexwise PhlowerTensor object. + original_pattern : str + Pattern of the original shape. Can be used for rearrange. + resultant_pattern : str + Pattern of the resultant shape. Can be used for rearrange. + dict_shape : dict[str, int] + Dict of original shape. Can be used for rearrange. + """ shape = self.shape dict_shape = {} @@ -220,14 +213,11 @@ def to_2d(self) -> PhlowerTensor: t_pattern = "" if self.is_voxel: space_pattern = "x y z " - try: - dict_shape.update({ - "x": shape[space_start], - "y": shape[space_start + 1], - "z": shape[space_start + 2], - }) - except: - raise ValueError(self.shape, self.is_time_series, self.is_voxel) + dict_shape.update({ + "x": shape[space_start], + "y": shape[space_start + 1], + "z": shape[space_start + 2], + }) feat_start = space_start + 3 else: space_pattern = "n " @@ -241,49 +231,23 @@ def to_2d(self) -> PhlowerTensor: f"a{i}": s for i, s in enumerate(shape[feat_start:-1])}) original_pattern = f"{t_pattern}{space_pattern}{feat_pattern}" - current_pattern = f"({space_pattern}) ({t_pattern}{feat_pattern})" + resultant_pattern = f"({space_pattern}) ({t_pattern}{feat_pattern})" tensor_2d = einops.rearrange( - self.to_tensor(), f"{original_pattern} -> {current_pattern}") + self.to_tensor(), f"{original_pattern} -> {resultant_pattern}") return PhlowerTensor( tensor_2d, dimension_tensor=self.dimension, - is_time_series=False, is_voxel=False, - original_pattern=original_pattern, - current_pattern=current_pattern, - dict_shape=dict_shape) - - def from_2d(self) -> PhlowerTensor: - if self.original_pattern is None: - raise PhlowerReshapeError( - "No original_pattern found. Run to_2d first.") - is_time_series = "t" in self.original_pattern - is_voxel = "x" in self.original_pattern - return self.inverse_rearrange( - is_time_series=is_time_series, is_voxel=is_voxel) + is_time_series=False, is_voxel=False), \ + original_pattern, resultant_pattern, dict_shape def rearrange( self, pattern: str, is_time_series: bool = False, is_voxel: bool = False, **kwargs) -> PhlowerTensor: tensor = self.to_tensor() - original_pattern, current_pattern = pattern.split("->") rearranged = einops.rearrange(tensor, pattern, **kwargs) return PhlowerTensor( rearranged, dimension_tensor=self.dimension, - is_time_series=is_time_series, is_voxel=is_voxel, - original_pattern=original_pattern, current_pattern=current_pattern, - dict_shape=kwargs) - - def inverse_rearrange( - self, is_time_series: bool = False, is_voxel: bool = False, - **kwargs) -> PhlowerTensor: - if self.original_pattern is None: - raise PhlowerReshapeError( - "No original_pattern found. Run rearrange first.") - pattern = f"{self.current_pattern} -> {self.original_pattern}" - kwargs.update(self.dict_shape) - return self.rearrange( - pattern, is_time_series=is_time_series, is_voxel=is_voxel, - **kwargs) + is_time_series=is_time_series, is_voxel=is_voxel) def reshape( self, shape: Sequence[int], diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 99ad326..cc572cd 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -4,10 +4,27 @@ from phlower._base.tensors._interface import IPhlowerTensor -def spmm(sparse: IPhlowerTensor, x: IPhlowerTensor) -> IPhlowerTensor: - h = x.to_2d() - ret = torch.sparse.mm(sparse, h) - pattern = f"{h.current_pattern} -> {h.original_pattern}" - return ret.rearrange( +def spmm( + sparse: IPhlowerTensor, x: IPhlowerTensor, + repeat: int = 1) -> IPhlowerTensor: + """ + Computes sparse matrix times dense tensor along with the vertex axis. + + Args: + sparse : IPhlowerTensor: + Sparse tensor. + x : IPhlowerTensor + Dense tensor. + repeat : int, optional + The number of repetitions for multiplication. The default is 1. + Returns: + IPhlowerTensor: + Resultant tensor. + """ + h, original_pattern, resultant_pattern, dict_shape = x.to_vertexwise() + pattern = f"{resultant_pattern} -> {original_pattern}" + for _ in range(repeat): + h = torch.sparse.mm(sparse, h) + return h.rearrange( pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, - **h.dict_shape) + **dict_shape) diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index f041dde..d0192c0 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -4,7 +4,7 @@ from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections -from phlower.nn._core_modules import _utils +from phlower.nn._core_modules import _functions, _utils from phlower.nn._interface_module import ( IPhlowerCoreModule, IReadonlyReferenceGroup, @@ -102,19 +102,5 @@ def forward( def _propagate( self, x: PhlowerTensor, support: PhlowerTensor ) -> PhlowerTensor: - # NOTE: Could be simplified as follows, - # but maybe slow due to reshape running every time - # from phlower.nn._core_modules import _utils - # h = x - # for _ in range(self._repeat): - # h = _functions.spmm(support, h) * self._factor - # return h - - h = x.to_2d() - pattern = f"{h.current_pattern} -> {h.original_pattern}" - dict_shape = h.dict_shape - for _ in range(self._repeat): - h = torch.sparse.mm(support, h) * self._factor - return h.rearrange( - pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, - **dict_shape) + h = _functions.spmm(support * self._factor, x, repeat=self._repeat) + return h diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index 2d54fe5..4292e55 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -158,11 +158,11 @@ def test__raises_phlower_sparse_rank_undefined_error(): ( True, True, [4, 10, 10, 10, 3, 3, 16], (1000, 4 * 3 * 3 * 16)), ], ) -def test__to_2d(is_time_series, is_voxel, size, desired_shape): +def test__to_vertexwise(is_time_series, is_voxel, size, desired_shape): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) - assert phlower_tensor.to_2d().shape == desired_shape + assert phlower_tensor.to_vertexwise()[0].shape == desired_shape @pytest.mark.parametrize( @@ -182,11 +182,17 @@ def test__to_2d(is_time_series, is_voxel, size, desired_shape): ( True, True, [4, 10, 10, 10, 3, 3, 16]), ], ) -def test__to_from_2d(is_time_series, is_voxel, size): +def test__to_vertexwise_inverse(is_time_series, is_voxel, size): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) - actual = phlower_tensor.to_2d().from_2d() + vertexwise, original_pattern, resultant_pattern, dict_shape\ + = phlower_tensor.to_vertexwise() + assert len(vertexwise.shape) == 2 + pattern = f"{resultant_pattern} -> {original_pattern}" + actual = vertexwise.rearrange( + pattern, is_time_series=is_time_series, is_voxel=is_voxel, + **dict_shape) np.testing.assert_almost_equal( actual.to_tensor().numpy(), phlower_tensor.to_tensor().numpy()) @@ -203,7 +209,3 @@ def test__rearrange(input_shape, pattern, dict_shape, desired_shape): phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) actual = phlower_tensor.rearrange(pattern, **dict_shape) assert actual.shape == desired_shape - assert actual.dict_shape == dict_shape - - actual_inversed = actual.inverse_rearrange() - assert actual_inversed.shape == phlower_tensor.shape diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 81ae3a9..f5ba46c 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -9,30 +9,39 @@ @pytest.mark.parametrize( - "size, is_time_series", + "size, is_time_series, repeat", [ - ((10, 1), False), - ((10, 16), False), - ((10, 3, 16), False), - ((4, 10, 1), True), - ((4, 10, 16), True), - ((4, 10, 3, 16), True), + ((10, 1), False, 1), + ((10, 16), False, 1), + ((10, 3, 16), False, 1), + ((4, 10, 1), True, 1), + ((4, 10, 16), True, 1), + ((4, 10, 3, 16), True, 1), + ((10, 1), False, 5), + ((10, 16), False, 5), + ((10, 3, 16), False, 5), + ((4, 10, 1), True, 5), + ((4, 10, 16), True, 5), + ((4, 10, 3, 16), True, 5), ], ) -def test__spmm(size, is_time_series): +def test__spmm(size, is_time_series, repeat): phlower_tensor = PhlowerTensor( torch.rand(*size), is_time_series=is_time_series) n = phlower_tensor.n_vertices() sparse = PhlowerTensor(torch.rand(n, n).to_sparse()) - actual_spmm = _functions.spmm(sparse, phlower_tensor).to_tensor().numpy() + actual_spmm = _functions.spmm( + sparse, phlower_tensor, repeat=repeat).to_tensor().numpy() sp_sparse = sp.coo_array(sparse.to_tensor().to_dense().numpy()) np_dense = phlower_tensor.to_tensor().numpy() def assert_correct(actual, array): dim_feat = len(array.shape) - 1 if dim_feat == 1: - desired = sp_sparse @ array + desired = array + for _ in range(repeat): + desired = sp_sparse @ desired np.testing.assert_almost_equal(actual, desired, decimal=5) return From 4a02c127e8661fbc21e6711c763a7af97097834f Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 00:18:25 +0900 Subject: [PATCH 22/89] add einsum and contraction --- .../_base/tensors/_dimension_tensor.py | 7 +- src/phlower/_base/tensors/_phlower_tensor.py | 7 + .../nn/_core_modules/_en_equivariant_mlp.py | 92 +++++++ src/phlower/nn/_core_modules/_functions.py | 106 ++++++++ .../settings/_module_settings/__init__.py | 4 + .../_en_equivariant_mlp_setting.py | 71 +++++ src/phlower/utils/enums.py | 2 +- src/phlower/utils/exceptions.py | 7 + .../test_core_modules/test_functions.py | 251 +++++++++++++++++- 9 files changed, 543 insertions(+), 4 deletions(-) create mode 100644 src/phlower/nn/_core_modules/_en_equivariant_mlp.py create mode 100644 src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 268cd79..b0eddb6 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -43,8 +43,8 @@ def from_list( Args: values (list[float] | tuple[float]): list or tuple - if length of values is not equal to the number of registered dimension type, - raise ValueError. + if length of values is not equal to the number of registered + dimension type, raise ValueError. Returns: PhlowerDimensionTensor: tensor object @@ -76,6 +76,9 @@ def __init__( def __add__(self, __value: object): return torch.add(self, __value) + def __mul__(self, __value: object): + return torch.mul(self, __value) + def __eq__(self, other: object) -> bool: if not isinstance(other, PhlowerDimensionTensor): return NotImplemented diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 73a2b6a..bf1f561 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -34,6 +34,8 @@ def phlower_tensor( | PhlowerDimensionTensor | torch.Tensor | dict[str, float] + | list[float] + | tuple[float] | None ) = None, is_time_series: bool = False, @@ -59,6 +61,8 @@ def _resolve_dimension_arg( | PhlowerDimensionTensor | torch.Tensor | dict[str, float] + | list[float] + | tuple[float] | None, ) -> PhlowerDimensionTensor | None: if inputs is None: @@ -73,6 +77,9 @@ def _resolve_dimension_arg( if isinstance(inputs, dict | PhysicalDimensions): return phlower_dimension_tensor(inputs) + if isinstance(inputs, list | tuple): + return PhlowerDimensionTensor.from_list(inputs) + raise NotImplementedError( f"{type(inputs)} is not implemented " "when creating PhlowerDimensionTensor" diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py new file mode 100644 index 0000000..ca1b4af --- /dev/null +++ b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import torch +from typing_extensions import Self + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._core_modules import _utils +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings import EnEquivariantMLPSetting + + +class EnEquivariantMLP(IPhlowerCoreModule, torch.nn.Module): + """E(n)-equivariant Multi Layer Perceptron""" + + @classmethod + def from_setting(cls, setting: EnEquivariantMLPSetting) -> Self: + """Generate model from setting object + + Args: + setting (EnEquivariantMLPSetting): setting object + + Returns: + Self: EnEquivariantMLP object + """ + return cls(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + """Return neural network name + + Returns: + str: name + """ + return "EnEquivariantMLP" + + @classmethod + def need_reference(cls) -> bool: + return False + + def __init__( + self, + nodes: list[int], + activations: list[str] | None = None, + dropouts: list[float] | None = None, + bias: bool = False, + ) -> None: + super().__init__() + + if activations is None: + activations = [] + if dropouts is None: + dropouts = [] + + self._chains = _utils.ExtendedLinearList( + nodes=nodes, activations=activations, dropouts=dropouts, bias=bias + ) + self._nodes = nodes + self._activations = activations + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + + def get_reference_name(self) -> str | None: + return None + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + **kwards, + ) -> PhlowerTensor: + """forward function which overloads torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor], optional): + Graph object. Defaults to None. + + Returns: + PhlowerTensor: Tensor object + """ + h = data.unique_item() + for i in range(len(self._chains)): + h = self._chains.forward(h, index=i) + return h diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index cc572cd..f19e511 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -1,7 +1,11 @@ import torch +from phlower._base._dimension import PhysicalDimensions +from phlower._base.tensors import phlower_tensor +from phlower._base.tensors._dimension_tensor import PhlowerDimensionTensor from phlower._base.tensors._interface import IPhlowerTensor +from phlower.utils.exceptions import PhlowerIncompatibleTensorError def spmm( @@ -28,3 +32,105 @@ def spmm( return h.rearrange( pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, **dict_shape) + + +def contraction( + x: IPhlowerTensor, y: IPhlowerTensor | None = None) -> IPhlowerTensor: + """ + Compute the tensor contraction. + + Args: + x : IPhlowerTensor + Input tensor. + y : IPhlowerTensor, optional + Second input tensor. If not fed, it computes contraction of x. + """ + if y is None: + y = x + + if x.rank() < y.rank(): + return contraction(y, x) + # Now we assume x.rank() >= y.rank() + + if x.is_voxel != y.is_voxel: + raise PhlowerIncompatibleTensorError( + "Cannot compute contraction between non-voxel and voxel.") + + ret_is_time_series = False + if x.is_time_series: + time_x = "t" + ret_is_time_series = True + else: + time_x = "" + if y.is_time_series: + time_y = "t" + ret_is_time_series = True + else: + time_y = "" + if ret_is_time_series: + time_ret = "t" + else: + time_ret = "" + + # No need to consider y because they should be compatible + if x.is_voxel: + space = "xyz" + is_voxel = True + else: + space = "x" + is_voxel = False + + diff_rank = x.rank() - y.rank() + unresolved = "abcdeghijklmnopqrsuvw"[ + :diff_rank] # No f, t, x, y, and z because they are used elsewhere + + if x.dimension is None or y.dimension is None: + dimension = None + else: + dimension = x.dimension * y.dimension + return einsum( + f"{time_x}{space}...{unresolved}f,{time_y}{space}...f->" + f"{time_ret}{space}{unresolved}f", + x, y, dimension=dimension, + is_time_series=ret_is_time_series, is_voxel=is_voxel) + + +def einsum( + equation, *args: list[IPhlowerTensor], + dimension: ( + PhysicalDimensions + | PhlowerDimensionTensor + | torch.Tensor + | dict[str, float] + | list[float] + | tuple[float] + | None + ) = None, + is_time_series: bool = False, is_voxel: bool = False, +) -> IPhlowerTensor: + """ + Compute einsum for phlower tensors. + + Args: + equation: str + Equation for einsum operation. + args: list[IPhlowerTensor] + List of IPhlowerTensor objects. + dimension: + PhlowerDimensions | PhlowerDimensionTensor | torch.Tensor + | dict[str, float] | list[float] | tuple[float] | None + Dimension for the resultant tensor. + is_time_series: bool, optional + Flag for time series. The default is False. + is_voxel: bool, optional + Flag for voxel. The default is False. + + Returns: + IPhlowerTensor: + Resultant tensor + """ + ret_tensor = torch.einsum( + equation, [a.to_tensor() for a in args]) + return phlower_tensor( + ret_tensor, dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index d576583..0e06b55 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -5,11 +5,15 @@ from phlower.settings._module_settings._concatenator_setting import ( ConcatenatorSetting, ) +from phlower.settings._module_settings._en_equivariant_mlp_setting import ( + EnEquivariantMLPSetting, +) from phlower.settings._module_settings._gcn_setting import GCNSetting from phlower.settings._module_settings._mlp_setting import MLPSetting from phlower.settings._module_settings._share_setting import ShareSetting _name_to_setting: dict[str, IPhlowerLayerParameters] = { + "EnEquivariantMLP": EnEquivariantMLPSetting, "GCN": GCNSetting, "MLP": MLPSetting, "Concatenator": ConcatenatorSetting, diff --git a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py new file mode 100644 index 0000000..bb50a11 --- /dev/null +++ b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import pydantic +from pydantic import Field +from typing_extensions import Self + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) + + +class EnEquivariantMLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): + nodes: list[int] = Field( + ... + ) # This property only overwritten when resolving. + activations: list[str] = Field(default_factory=lambda: [], frozen=True) + dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) + bias: bool = Field(False, frozen=True) + + def gather_input_dims(self, *input_dims: int) -> int: + if len(input_dims) != 1: + raise ValueError("Only one input is allowed in EnEquivariantMLP.") + return input_dims[0] + + @pydantic.field_validator("nodes") + @classmethod + def check_n_nodes(cls, vals: list[int]) -> list[int]: + pass + if len(vals) < 2: + raise ValueError( + "size of nodes must be larger than 1 in " + "EnEquivariantMLPSetting. input: {vals}" + ) + + for i, v in enumerate(vals): + if v > 0: + continue + + if (i == 0) and (v == -1): + continue + + raise ValueError( + "nodes in EnEquivariantMLPSetting is inconsistent. " + f"value {v} in {i}-th of nodes is not allowed." + ) + + return vals + + @pydantic.model_validator(mode="after") + def check_nodes_size(self) -> Self: + if len(self.nodes) - 1 != len(self.activations): + raise ValueError( + "Size of nodes and activations is not compatible " + "in EnEquivariantMLPSetting." + " len(nodes) must be equal to 1 + len(activations)." + ) + return self + + def get_n_nodes(self) -> list[int] | None: + return self.nodes + + def overwrite_nodes(self, nodes: list[int]) -> None: + self.nodes = nodes + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return diff --git a/src/phlower/utils/enums.py b/src/phlower/utils/enums.py index 78a5d32..48b3e39 100644 --- a/src/phlower/utils/enums.py +++ b/src/phlower/utils/enums.py @@ -52,7 +52,7 @@ class PhysicalDimensionSymbolType(Enum): T = 0 # time L = 1 # length M = 2 # mass - I = 3 # electric current + I = 3 # electric current # NOQA Theta = 4 # thermodynamic temperature N = 5 # amount of substance J = 6 # luminous intensity diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 9dbe5f6..78c61eb 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -64,3 +64,10 @@ class PhlowerReshapeError(ValueError): """ class NotFoundReferenceModuleError(ValueError): ... + + +class PhlowerIncompatibleTensorError(ValueError): + """ + This error raises when trying to perform an operation for incompatible + tensor(s) + """ diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index f5ba46c..903e361 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -4,8 +4,9 @@ import torch from scipy import sparse as sp -from phlower import PhlowerTensor +from phlower import PhlowerTensor, phlower_tensor from phlower.nn._core_modules import _functions +from phlower.utils.exceptions import PhlowerIncompatibleTensorError @pytest.mark.parametrize( @@ -54,3 +55,251 @@ def assert_correct(actual, array): assert_correct(actual_spmm[t], np_dense[t]) else: assert_correct(actual_spmm, np_dense) + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel, desired_pattern", + [ + ((10, 1), False, False, "n...f,n...f->nf"), + ((10, 16), False, False, "n...f,n...f->nf"), + ((10, 3, 16), False, False, "n...f,n...f->nf"), + ((10, 3, 3, 16), False, False, "n...f,n...f->nf"), + ((4, 10, 1), True, False, "tn...f,tn...f->tnf"), + ((4, 10, 16), True, False, "tn...f,tn...f->tnf"), + ((4, 10, 3, 16), True, False, "tn...f,tn...f->tnf"), + ((4, 10, 3, 3, 16), True, False, "tn...f,tn...f->tnf"), + ((10, 10, 10, 1), False, True, "xyz...f,xyz...f->xyzf"), + ((10, 10, 10, 16), False, True, "xyz...f,xyz...f->xyzf"), + ((10, 10, 10, 3, 16), False, True, "xyz...f,xyz...f->xyzf"), + ((10, 10, 10, 3, 3, 16), False, True, "xyz...f,xyz...f->xyzf"), + ((4, 10, 10, 10, 1), True, True, "txyz...f,txyz...f->txyzf"), + ((4, 10, 10, 10, 16), True, True, "txyz...f,txyz...f->txyzf"), + ((4, 10, 10, 10, 3, 16), True, True, "txyz...f,txyz...f->txyzf"), + ((4, 10, 10, 10, 3, 3, 16), True, True, "txyz...f,txyz...f->txyzf"), + ], +) +@pytest.mark.parametrize( + "dimension", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +def test_contraction_one_argument_non_timeseries_non_voxel( + size, is_time_series, is_voxel, desired_pattern, dimension): + torch_tensor = torch.rand(*size) + x = phlower_tensor( + torch_tensor, dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + actual = _functions.contraction(x) + desired = torch.einsum( + desired_pattern, torch_tensor, torch_tensor).numpy() + np.testing.assert_almost_equal(actual.to_tensor().numpy(), desired) + + assert actual.is_time_series == is_time_series + assert actual.is_voxel == is_voxel + assert actual.rank() == 0 + + if dimension is not None: + actual_dimension = actual.dimension._tensor.numpy() + desired_dimension = np.array(dimension) * 2 + np.testing.assert_almost_equal(actual_dimension, desired_dimension) + else: + assert actual.dimension is None + + +@pytest.mark.parametrize( + "size_x, size_y, x_is_time_series, y_is_time_series, is_voxel, " + "desired_pattern, desired_rank", + [ + # Base + ( + (10, 16), (10, 16), False, False, False, + "nf,nf->nf", 0), + ( + (10, 3, 16), (10, 16), False, False, False, + "npf,nf->npf", 1), + ( + (10, 16), (10, 3, 16), False, False, False, + "nf,npf->npf", 1), + ( + (10, 3, 16), (10, 3, 16), False, False, False, + "npf,npf->nf", 0), + ( + (10, 3, 3, 16), (10, 16), False, False, False, + "npqf,nf->npqf", 2), + ( + (10, 3, 3, 16), (10, 3, 16), False, False, False, + "npqf,npf->nqf", 1), + ( + (10, 3, 3, 16), (10, 3, 3, 16), False, False, False, + "npqf,npqf->nf", 0), + # X time series + ( + (4, 10, 16), (10, 16), True, False, False, + "tnf,nf->tnf", 0), + ( + (4, 10, 3, 16), (10, 16), True, False, False, + "tnpf,nf->tnpf", 1), + ( + (4, 10, 16), (10, 3, 16), True, False, False, + "tnf,npf->tnpf", 1), + ( + (4, 10, 3, 16), (10, 3, 16), True, False, False, + "tnpf,npf->tnf", 0), + ( + (4, 10, 3, 3, 16), (10, 16), True, False, False, + "tnpqf,nf->tnpqf", 2), + ( + (4, 10, 3, 3, 16), (10, 3, 16), True, False, False, + "tnpqf,npf->tnqf", 1), + ( + (4, 10, 3, 3, 16), (10, 3, 3, 16), True, False, False, + "tnpqf,npqf->tnf", 0), + # Y time series + ( + (10, 16), (4, 10, 16), False, True, False, + "nf,tnf->tnf", 0), + ( + (10, 3, 16), (4, 10, 16), False, True, False, + "npf,tnf->tnpf", 1), + ( + (10, 16), (4, 10, 3, 16), False, True, False, + "nf,tnpf->tnpf", 1), + ( + (10, 3, 16), (4, 10, 3, 16), False, True, False, + "npf,tnpf->tnf", 0), + ( + (10, 3, 3, 16), (4, 10, 16), False, True, False, + "npqf,tnf->tnpqf", 2), + ( + (10, 3, 3, 16), (4, 10, 3, 16), False, True, False, + "npqf,tnpf->tnqf", 1), + ( + (10, 3, 3, 16), (4, 10, 3, 3, 16), False, True, False, + "npqf,tnpqf->tnf", 0), + # X Y time series + ( + (4, 10, 16), (4, 10, 16), True, True, False, + "tnf,tnf->tnf", 0), + ( + (4, 10, 3, 16), (4, 10, 16), True, True, False, + "tnpf,tnf->tnpf", 1), + ( + (4, 10, 16), (4, 10, 3, 16), True, True, False, + "tnf,tnpf->tnpf", 1), + ( + (4, 10, 3, 16), (4, 10, 3, 16), True, True, False, + "tnpf,tnpf->tnf", 0), + ( + (4, 10, 3, 3, 16), (4, 10, 16), True, True, False, + "tnpqf,tnf->tnpqf", 2), + ( + (4, 10, 3, 3, 16), (4, 10, 3, 16), True, True, False, + "tnpqf,tnpf->tnqf", 1), + ( + (4, 10, 3, 3, 16), (4, 10, 3, 3, 16), True, True, False, + "tnpqf,tnpqf->tnf", 0), + # X time series, X Y voxel + ( + (4, 10, 10, 10, 16), (10, 10, 10, 16), + True, False, True, "txyzf,xyzf->txyzf", 0), + ( + (4, 10, 10, 10, 3, 16), (10, 10, 10, 16), + True, False, True, "txyzpf,xyzf->txyzpf", 1), + ( + (4, 10, 10, 10, 16), (10, 10, 10, 3, 16), + True, False, True, "txyzf,xyzpf->txyzpf", 1), + ( + (4, 10, 10, 10, 3, 16), (10, 10, 10, 3, 16), + True, False, True, "txyzpf,xyzpf->txyzf", 0), + ( + (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 16), + True, False, True, "txyzpqf,xyzf->txyzpqf", 2), + ( + (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 16), + True, False, True, "txyzpqf,xyzpf->txyzqf", 1), + ( + (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 3, 16), + True, False, True, "txyzpqf,xyzpqf->txyzf", 0), + # X Y time series, X Y voxel + ( + (4, 10, 10, 10, 16), (4, 10, 10, 10, 16), + True, True, True, "txyzf,txyzf->txyzf", 0), + ( + (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 16), + True, True, True, "txyzpf,txyzf->txyzpf", 1), + ( + (4, 10, 10, 10, 16), (4, 10, 10, 10, 3, 16), + True, True, True, "txyzf,txyzpf->txyzpf", 1), + ( + (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 3, 16), + True, True, True, "txyzpf,txyzpf->txyzf", 0), + ( + (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 16), + True, True, True, "txyzpqf,txyzf->txyzpqf", 2), + ( + (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 16), + True, True, True, "txyzpqf,txyzpf->txyzqf", 1), + ( + (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 3, 16), + True, True, True, "txyzpqf,txyzpqf->txyzf", 0), + ], +) +@pytest.mark.parametrize( + "dimension_x", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +@pytest.mark.parametrize( + "dimension_y", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +def test_contraction_two_arguments( + size_x, size_y, x_is_time_series, y_is_time_series, is_voxel, + desired_pattern, desired_rank, dimension_x, dimension_y): + t_x = torch.rand(*size_x) + x = phlower_tensor( + t_x, dimension=dimension_x, + is_time_series=x_is_time_series, is_voxel=is_voxel) + + t_y = torch.rand(*size_y) + y = phlower_tensor( + t_y, dimension=dimension_y, + is_time_series=y_is_time_series, is_voxel=is_voxel) + + actual = _functions.contraction(x, y) + desired = torch.einsum( + desired_pattern, t_x, t_y).numpy() + np.testing.assert_almost_equal(actual.to_tensor().numpy(), desired) + + assert actual.is_time_series == x_is_time_series or y_is_time_series + assert actual.is_voxel == is_voxel + assert actual.rank() == desired_rank + + if dimension_x is not None and dimension_y is not None: + actual_dimension = actual.dimension._tensor.numpy() + desired_dimension = np.array(dimension_x) + np.array(dimension_y) + np.testing.assert_almost_equal(actual_dimension, desired_dimension) + else: + assert actual.dimension is None + + +def test_contraction_raises_phlower_incompatible_tensor_error(): + x = phlower_tensor( + torch.rand(10, 10, 10, 3, 16), is_voxel=True) + y = phlower_tensor( + torch.rand(10 * 10 * 10, 3, 16), is_voxel=False) + with pytest.raises(PhlowerIncompatibleTensorError): + _functions.contraction(x, y) From 66c4d462c7119f522502b2caa43d7cb2eb8587fc Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 03:31:55 +0900 Subject: [PATCH 23/89] add E(n)-Equivariant MLP and necessary operations --- src/phlower/_base/tensors/_phlower_tensor.py | 8 +- src/phlower/nn/__init__.py | 6 + src/phlower/nn/_core_modules/__init__.py | 16 +- .../nn/_core_modules/_en_equivariant_mlp.py | 31 ++- src/phlower/nn/_core_modules/_functions.py | 186 ++++++++++--- src/phlower/nn/_core_modules/_gcn.py | 2 +- src/phlower/nn/_core_modules/_identity.py | 72 +++++ src/phlower/nn/_core_modules/_mlp.py | 3 +- src/phlower/nn/_core_modules/_proportional.py | 85 ++++++ src/phlower/nn/_core_modules/_utils.py | 25 +- .../settings/_module_settings/__init__.py | 8 +- .../_en_equivariant_mlp_setting.py | 4 +- .../_module_settings/_identity_setting.py | 50 ++++ .../_module_settings/_proportional_setting.py | 62 +++++ src/phlower/utils/exceptions.py | 3 + .../test_en_equivariant_mlp.py | 58 ++++ .../test_core_modules/test_functions.py | 254 +++++++++++++++++- .../test_core_modules/test_identity.py | 34 +++ .../test_core_modules/test_proportional.py | 45 ++++ 19 files changed, 895 insertions(+), 57 deletions(-) create mode 100644 src/phlower/nn/_core_modules/_identity.py create mode 100644 src/phlower/nn/_core_modules/_proportional.py create mode 100644 src/phlower/settings/_module_settings/_identity_setting.py create mode 100644 src/phlower/settings/_module_settings/_proportional_setting.py create mode 100644 tests/test_nn/test_core_modules/test_en_equivariant_mlp.py create mode 100644 tests/test_nn/test_core_modules/test_identity.py create mode 100644 tests/test_nn/test_core_modules/test_proportional.py diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index bf1f561..84e1a07 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -21,6 +21,7 @@ from phlower.utils.exceptions import ( DimensionIncompatibleError, PhlowerSparseUnsupportedError, + PhlowerTypeError, PhlowerUnsupportedTorchFunctionError, ) @@ -100,7 +101,9 @@ def __init__( is_time_series: bool = False, is_voxel: bool = False, ): - assert isinstance(tensor, torch.Tensor) + if not isinstance(tensor, torch.Tensor): + raise PhlowerTypeError( + f"Expect torch.Tensor but {tensor.__class__} was fed") self._tensor = tensor self._dimension_tensor = dimension_tensor self._is_time_series = is_time_series @@ -154,6 +157,9 @@ def __len__(self) -> int: def to_tensor(self) -> torch.Tensor: return self._tensor + def to_numpy(self) -> np.ndarray: + return self._tensor.cpu().detach().numpy() + def coalesce(self) -> torch.Tensor: return PhlowerTensor(self._tensor.coalesce(), self._dimension_tensor) diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 1af9012..17f1131 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -1,6 +1,12 @@ from phlower.nn._core_modules._concatenator import Concatenator from phlower.nn._core_modules._gcn import GCN +from phlower.nn._core_modules._identity import Identity from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share from phlower.nn._group_module import PhlowerGroupModule from phlower.nn._interface_module import IPhlowerCoreModule + +if True: + # NOTE: Import advanced models after + from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP diff --git a/src/phlower/nn/_core_modules/__init__.py b/src/phlower/nn/_core_modules/__init__.py index 17800d6..65ecce9 100644 --- a/src/phlower/nn/_core_modules/__init__.py +++ b/src/phlower/nn/_core_modules/__init__.py @@ -1,10 +1,24 @@ from phlower.nn._core_modules._concatenator import Concatenator from phlower.nn._core_modules._gcn import GCN +from phlower.nn._core_modules._identity import Identity from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share from phlower.nn._interface_module import IPhlowerCoreModule -_all_models: list[IPhlowerCoreModule] = [GCN, MLP, Concatenator, Share] +if True: + # NOTE: Import advanced models after + from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP + +_all_models: list[IPhlowerCoreModule] = [ + Concatenator, + EnEquivariantMLP, + GCN, + Identity, + MLP, + Proportional, + Share, +] _name2model = {cls.get_nn_name(): cls for cls in _all_models} diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py index ca1b4af..96aa9e8 100644 --- a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py @@ -5,7 +5,7 @@ from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections -from phlower.nn._core_modules import _utils +from phlower.nn._core_modules import Identity, Proportional, _functions, _utils from phlower.nn._interface_module import ( IPhlowerCoreModule, IReadonlyReferenceGroup, @@ -47,6 +47,8 @@ def __init__( activations: list[str] | None = None, dropouts: list[float] | None = None, bias: bool = False, + create_linear_weight: bool = False, + norm_function_name: str = None, ) -> None: super().__init__() @@ -56,10 +58,25 @@ def __init__( dropouts = [] self._chains = _utils.ExtendedLinearList( - nodes=nodes, activations=activations, dropouts=dropouts, bias=bias - ) + nodes=nodes, activations=activations, dropouts=dropouts, bias=bias) self._nodes = nodes self._activations = activations + self._create_linear_weight = create_linear_weight + self._norm_function_name = norm_function_name + + self._linear_weight = self._init_linear_weight() + self._norm_function = _utils.ActivationSelector.select( + self._norm_function_name) + + def _init_linear_weight(self): + if not self._create_linear_weight: + if self._nodes[0] != self._nodes[-1]: + raise ValueError( + "First and last nodes are different. " + "Set create_linear_weight True.") + return Identity() + + return Proportional([self._nodes[0], self._nodes[-1]]) def resolve( self, *, parent: IReadonlyReferenceGroup | None = None, **kwards @@ -86,7 +103,7 @@ def forward( Returns: PhlowerTensor: Tensor object """ - h = data.unique_item() - for i in range(len(self._chains)): - h = self._chains.forward(h, index=i) - return h + x = data.unique_item() + linear_x = self._linear_weight(data) + h = self._chains(self._norm_function(_functions.contraction(x))) + return _functions.tensor_product(linear_x, h) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index f19e511..7088c37 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -8,6 +8,10 @@ from phlower.utils.exceptions import PhlowerIncompatibleTensorError +def identity(x): + return x + + def spmm( sparse: IPhlowerTensor, x: IPhlowerTensor, repeat: int = 1) -> IPhlowerTensor: @@ -34,6 +38,61 @@ def spmm( **dict_shape) +def einsum( + equation, *args: list[IPhlowerTensor], + dimension: ( + PhysicalDimensions + | PhlowerDimensionTensor + | torch.Tensor + | dict[str, float] + | list[float] + | tuple[float] + | None + ) = None, + is_time_series: bool = False, is_voxel: bool = False, +) -> IPhlowerTensor: + """ + Compute einsum for phlower tensors. + + Args: + equation: str + Equation for einsum operation. + args: list[IPhlowerTensor] + List of IPhlowerTensor objects. + dimension: + PhlowerDimensions | PhlowerDimensionTensor | torch.Tensor + | dict[str, float] | list[float] | tuple[float] | None + Dimension for the resultant tensor. + is_time_series: bool, optional + Flag for time series. The default is False. + is_voxel: bool, optional + Flag for voxel. The default is False. + + Returns: + IPhlowerTensor: + Resultant tensor + """ + try: + ret_tensor = torch.einsum( + equation, [a.to_tensor() for a in args]) + except RuntimeError as e: + raise PhlowerIncompatibleTensorError( + f"{e}\n" + f"{equation}, {[a.shape for a in args]}") from e + return phlower_tensor( + ret_tensor, dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + + +def _availale_variables(length: int, start: int = 0) -> str: + # No f, t, x, y, and z because they are "reserved" + available_variables = "abcdeghijklmnopqrsuvw" + + if length > len(available_variables): + raise ValueError(f"Required length too long: {length}") + return available_variables[start:start+length] + + def contraction( x: IPhlowerTensor, y: IPhlowerTensor | None = None) -> IPhlowerTensor: """ @@ -81,8 +140,7 @@ def contraction( is_voxel = False diff_rank = x.rank() - y.rank() - unresolved = "abcdeghijklmnopqrsuvw"[ - :diff_rank] # No f, t, x, y, and z because they are used elsewhere + unresolved = _availale_variables(diff_rank) if x.dimension is None or y.dimension is None: dimension = None @@ -95,42 +153,104 @@ def contraction( is_time_series=ret_is_time_series, is_voxel=is_voxel) -def einsum( - equation, *args: list[IPhlowerTensor], - dimension: ( - PhysicalDimensions - | PhlowerDimensionTensor - | torch.Tensor - | dict[str, float] - | list[float] - | tuple[float] - | None - ) = None, - is_time_series: bool = False, is_voxel: bool = False, +def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: + """ + Compute spatial tensor product for input tensors. + + Args: + x: IPhlowerTensor + y: IPhlowerTensor + + Returns: + IPhlowerTensor: + Resultant tensor. + """ + if x.is_voxel != y.is_voxel: + raise PhlowerIncompatibleTensorError( + "Cannot compute contraction between non-voxel and voxel.") + + x_rank = x.rank() + y_rank = y.rank() + + ret_is_time_series = False + if x.is_time_series: + x_time = "t" + ret_is_time_series = True + else: + x_time = "" + if y.is_time_series: + y_time = "t" + ret_is_time_series = True + else: + y_time = "" + if ret_is_time_series: + t_ret = "t" + else: + t_ret = "" + + # No need to consider y because they should be compatible + if x.is_voxel: + space = "xyz" + is_voxel = True + else: + space = "x" + is_voxel = False + + x_vars = _availale_variables(x_rank) + y_vars = _availale_variables(y_rank, start=x_rank) + equation = f"{x_time}{space}{x_vars}f,{y_time}{space}{y_vars}f->" \ + + f"{t_ret}{space}{x_vars}{y_vars}f" + + if x.dimension is None or y.dimension is None: + dimension = None + else: + dimension = x.dimension * y.dimension + + return einsum( + equation, x, y, dimension=dimension, + is_time_series=ret_is_time_series, is_voxel=is_voxel) + + +def apply_orthogonal_group( + orthogonal_matrix: IPhlowerTensor, tensor: IPhlowerTensor ) -> IPhlowerTensor: """ - Compute einsum for phlower tensors. + Apply orthogonal group action to the input tensor. Args: - equation: str - Equation for einsum operation. - args: list[IPhlowerTensor] - List of IPhlowerTensor objects. - dimension: - PhlowerDimensions | PhlowerDimensionTensor | torch.Tensor - | dict[str, float] | list[float] | tuple[float] | None - Dimension for the resultant tensor. - is_time_series: bool, optional - Flag for time series. The default is False. - is_voxel: bool, optional - Flag for voxel. The default is False. + orthogonal_matrix: IPhlowerTensor + [3, 3]-shaped orthogonal matrix. + tensor: IPhlowerTensor + Tensor to apply the orthogonal group action. Returns: IPhlowerTensor: - Resultant tensor + Resultant tensor. """ - ret_tensor = torch.einsum( - equation, [a.to_tensor() for a in args]) - return phlower_tensor( - ret_tensor, dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + rank = tensor.rank() + if rank == 0: + return tensor + + start_dim = 1 + if tensor.is_time_series: + time = "t" + start_dim += 1 + else: + time = "" + if tensor.is_voxel: + space = "xyz" + start_dim += 3 + else: + space = "x" + + s = _availale_variables(rank * 2) + str_ortho = ','.join(a + b for a, b in zip(s[::2], s[1::2], strict=True)) + str_tensor = f"{time}{space}{s[1::2]}f" + str_ret = f"{time}{space}{s[::2]}f" + equation = f"{str_ortho},{str_tensor}->{str_ret}" + args = [orthogonal_matrix] * rank + [tensor] + + return einsum( + equation, *args, + dimension=tensor.dimension, + is_time_series=tensor.is_time_series, is_voxel=tensor.is_voxel) diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index d0192c0..b7a25b2 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -96,7 +96,7 @@ def forward( h = data.unique_item() for i in range(len(self._chains)): h = self._propagate(h, support) - h = self._chains.forward(h, index=i) + h = self._chains.forward_part(h, index=i) return h def _propagate( diff --git a/src/phlower/nn/_core_modules/_identity.py b/src/phlower/nn/_core_modules/_identity.py new file mode 100644 index 0000000..3bd356e --- /dev/null +++ b/src/phlower/nn/_core_modules/_identity.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch +from typing_extensions import Self + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings import IdentitySetting + + +class Identity(IPhlowerCoreModule, torch.nn.Module): + """Identity layer""" + + @classmethod + def from_setting(cls, setting: IdentitySetting) -> Self: + """Generate model from setting object + + Args: + setting (EnEquivariantMLPSetting): setting object + + Returns: + Self: Identity object + """ + return cls(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + """Return neural network name + + Returns: + str: name + """ + return "Identity" + + @classmethod + def need_reference(cls) -> bool: + return False + + def __init__(self, nodes: list[int] = None) -> None: + super().__init__() + self._nodes = nodes + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + + def get_reference_name(self) -> str | None: + return None + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + **kwards, + ) -> PhlowerTensor: + """forward function which overloads torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor], optional): + Graph object. Defaults to None. + + Returns: + PhlowerTensor: Tensor object + """ + return data.unique_item() diff --git a/src/phlower/nn/_core_modules/_mlp.py b/src/phlower/nn/_core_modules/_mlp.py index 60f28e5..90397e9 100644 --- a/src/phlower/nn/_core_modules/_mlp.py +++ b/src/phlower/nn/_core_modules/_mlp.py @@ -87,6 +87,5 @@ def forward( PhlowerTensor: Tensor object """ h = data.unique_item() - for i in range(len(self._chains)): - h = self._chains.forward(h, index=i) + h = self._chains.forward(h) return h diff --git a/src/phlower/nn/_core_modules/_proportional.py b/src/phlower/nn/_core_modules/_proportional.py new file mode 100644 index 0000000..35e211e --- /dev/null +++ b/src/phlower/nn/_core_modules/_proportional.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import torch +from typing_extensions import Self + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._core_modules import _utils +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings import ProportionalSetting + + +class Proportional(IPhlowerCoreModule, torch.nn.Module): + """Proportional, i.e., strictly linear, layer""" + + @classmethod + def from_setting(cls, setting: ProportionalSetting) -> Self: + """Generate model from setting object + + Args: + setting (ProportionalSetting): setting object + + Returns: + Self: Proportional object + """ + return cls(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + """Return neural network name + + Returns: + str: name + """ + return "Proportional" + + @classmethod + def need_reference(cls) -> bool: + return False + + def __init__( + self, + nodes: list[int], + dropouts: list[float] | None = None, + ) -> None: + super().__init__() + + if dropouts is None: + dropouts = [] + + self._chains = _utils.ExtendedLinearList( + nodes=nodes, activations=['identity'], + dropouts=[], bias=False) + self._nodes = nodes + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + + def get_reference_name(self) -> str | None: + return None + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + **kwards, + ) -> PhlowerTensor: + """forward function which overloads torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor], optional): + Graph object. Defaults to None. + + Returns: + PhlowerTensor: Tensor object + """ + h = data.unique_item() + return self._chains(h) diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index 55a16e8..2a038d4 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -3,6 +3,7 @@ import torch from phlower._base.tensors import PhlowerTensor +from phlower.nn._core_modules import _functions class ExtendedLinearList(torch.nn.Module): @@ -10,13 +11,15 @@ def __init__( self, nodes: list[int], activations: list[str], - dropouts: list[float], bias: bool, + dropouts: list[float] | None = None, ) -> None: super().__init__() self._nodes = nodes self._activations = activations + if dropouts is None: + dropouts = [] self._dropouts = dropouts self._validate_args() @@ -39,7 +42,7 @@ def __init__( def __len__(self) -> int: return len(self._linears) - def forward(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: + def forward_part(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: assert index < self._n_chains h = self._linears[index](x) @@ -49,6 +52,12 @@ def forward(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: h = self._activators[index](h) return h + def forward(self, x: PhlowerTensor) -> PhlowerTensor: + h = x + for i in range(self._n_chains - 1): + h = self.forward_part(h, index=i) + return h + def _validate_args(self) -> None: assert len(self._nodes) >= 2 @@ -65,19 +74,19 @@ def _validate_args(self) -> None: assert len(self._nodes) == len(self._dropouts) + 1 -def identity(x): - return x - - class ActivationSelector: _REGISTERED_ACTIVATIONS = { - "identity": identity, + "identity": _functions.identity, "relu": torch.relu, + "sigmoid": torch.sigmoid, + "sqrt": torch.sqrt, "tanh": torch.tanh, } @staticmethod - def select(name: str) -> Callable[[torch.Tensor], torch.Tensor]: + def select(name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: + if name is None: + name = "identity" return ActivationSelector._REGISTERED_ACTIVATIONS[name] @staticmethod diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index 0e06b55..73e351f 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -9,14 +9,20 @@ EnEquivariantMLPSetting, ) from phlower.settings._module_settings._gcn_setting import GCNSetting +from phlower.settings._module_settings._identity_setting import IdentitySetting from phlower.settings._module_settings._mlp_setting import MLPSetting +from phlower.settings._module_settings._proportional_setting import ( + ProportionalSetting, +) from phlower.settings._module_settings._share_setting import ShareSetting _name_to_setting: dict[str, IPhlowerLayerParameters] = { + "Concatenator": ConcatenatorSetting, "EnEquivariantMLP": EnEquivariantMLPSetting, "GCN": GCNSetting, + "Identity": IdentitySetting, "MLP": MLPSetting, - "Concatenator": ConcatenatorSetting, + "Proportional": ProportionalSetting, "Share": ShareSetting, } diff --git a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py index bb50a11..9774702 100644 --- a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py @@ -17,6 +17,9 @@ class EnEquivariantMLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): activations: list[str] = Field(default_factory=lambda: [], frozen=True) dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) bias: bool = Field(False, frozen=True) + create_linear_weight: bool = Field(False, frozen=True) + norm_function_name: str = Field( + default_factory=lambda: 'identity', frozen=True) def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: @@ -26,7 +29,6 @@ def gather_input_dims(self, *input_dims: int) -> int: @pydantic.field_validator("nodes") @classmethod def check_n_nodes(cls, vals: list[int]) -> list[int]: - pass if len(vals) < 2: raise ValueError( "size of nodes must be larger than 1 in " diff --git a/src/phlower/settings/_module_settings/_identity_setting.py b/src/phlower/settings/_module_settings/_identity_setting.py new file mode 100644 index 0000000..71cf6e7 --- /dev/null +++ b/src/phlower/settings/_module_settings/_identity_setting.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import pydantic +from pydantic import Field + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) + + +class IdentitySetting(IPhlowerLayerParameters, pydantic.BaseModel): + nodes: list[int] | None = Field( + None + ) # This property only overwritten when resolving. + activation: str = Field("identity", frozen=True) + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(extra="forbid") + + def gather_input_dims(self, *input_dims: int) -> int: + assert len(input_dims) > 0 + sum_dim = sum(v for v in input_dims) + return sum_dim + + @pydantic.field_validator("nodes") + @classmethod + def check_n_nodes(cls, vals: list[int]) -> list[int]: + if vals is None: + return vals + + if len(vals) != 2: + raise ValueError( + "size of nodes must be 2 in ConcatenatorSettings." + f" input: {vals}" + ) + return vals + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return + + def get_n_nodes(self) -> list[int] | None: + return self.nodes + + def overwrite_nodes(self, nodes: list[int]) -> None: + self.nodes = nodes diff --git a/src/phlower/settings/_module_settings/_proportional_setting.py b/src/phlower/settings/_module_settings/_proportional_setting.py new file mode 100644 index 0000000..a4ca1fa --- /dev/null +++ b/src/phlower/settings/_module_settings/_proportional_setting.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import pydantic +from pydantic import Field +from typing_extensions import Self + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) + + +class ProportionalSetting(IPhlowerLayerParameters, pydantic.BaseModel): + nodes: list[int] = Field( + ... + ) # This property only overwritten when resolving. + dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) + + def gather_input_dims(self, *input_dims: int) -> int: + if len(input_dims) != 1: + raise ValueError("Only one input is allowed in EnEquivariantMLP.") + return input_dims[0] + + @pydantic.field_validator("nodes") + @classmethod + def check_n_nodes(cls, vals: list[int]) -> list[int]: + if len(vals) == 2: + raise ValueError( + "size of nodes must be larger than 1 in " + "EnEquivariantMLPSetting. input: {vals}" + ) + + for i, v in enumerate(vals): + if v > 0: + continue + + if (i == 0) and (v == -1): + continue + + raise ValueError( + "nodes in ProportionalSetting is inconsistent. " + f"value {v} in {i}-th of nodes is not allowed." + ) + + return vals + + @pydantic.model_validator(mode="after") + def check_nodes_size(self) -> Self: + pass + + def get_n_nodes(self) -> list[int] | None: + return self.nodes + + def overwrite_nodes(self, nodes: list[int]) -> None: + self.nodes = nodes + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 78c61eb..11ea81b 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -71,3 +71,6 @@ class PhlowerIncompatibleTensorError(ValueError): This error raises when trying to perform an operation for incompatible tensor(s) """ + +class PhlowerTypeError(TypeError): + """This error raises when type is not compatible with what is expected.""" diff --git a/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py new file mode 100644 index 0000000..ffae0f3 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py @@ -0,0 +1,58 @@ +import numpy as np +import pytest +import torch +from scipy.stats import ortho_group + +from phlower import PhlowerTensor +from phlower.collections import phlower_tensor_collection +from phlower.nn import EnEquivariantMLP +from phlower.nn._core_modules import _functions + + +def test__can_call_parameters(): + model = EnEquivariantMLP(nodes=[8, 8]) + + # To check Concatenator inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel", + [ + ((10, 1), False, False), + ((10, 16), False, False), + ((10, 3, 16), False, False), + ((4, 10, 1), True, False), + ((4, 10, 16), True, False), + ((4, 10, 3, 16), True, False), + ((10, 10, 10, 1), False, True), + ((10, 10, 10, 16), False, True), + ((10, 10, 10, 3, 16), False, True), + ((4, 10, 10, 10, 1), True, True), + ((4, 10, 10, 10, 16), True, True), + ((4, 10, 10, 10, 3, 16), True, True), + ], +) +@pytest.mark.parametrize("n_output_feature", [1, 16, 32]) +def test__en_equivariance( + size, is_time_series, is_voxel, n_output_feature): + orthogonal_tensor = PhlowerTensor( + torch.tensor(ortho_group.rvs(3).astype(np.float32))) + create_linear_weight = size[-1] != n_output_feature + model = EnEquivariantMLP( + nodes=[size[-1], n_output_feature], + create_linear_weight=create_linear_weight) + + phlower_tensor = PhlowerTensor( + torch.rand(*size), is_time_series=is_time_series, is_voxel=is_voxel) + + phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) + actual = _functions.apply_orthogonal_group( + orthogonal_tensor, model(phlower_tensors)).to_numpy() + + rotated_phlower_tensors = phlower_tensor_collection( + {'tensor': _functions.apply_orthogonal_group( + orthogonal_tensor, phlower_tensor)}) + desired = model(rotated_phlower_tensors).to_numpy() + + np.testing.assert_almost_equal(actual, desired, decimal=6) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 903e361..cb6c062 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -3,6 +3,7 @@ import pytest import torch from scipy import sparse as sp +from scipy.stats import ortho_group from phlower import PhlowerTensor, phlower_tensor from phlower.nn._core_modules import _functions @@ -43,7 +44,9 @@ def assert_correct(actual, array): desired = array for _ in range(repeat): desired = sp_sparse @ desired - np.testing.assert_almost_equal(actual, desired, decimal=5) + norm = np.mean(np.linalg.norm(desired, axis=-1)) + np.testing.assert_almost_equal( + actual / norm, desired / norm, decimal=5) return for i in range(array.shape[1]): @@ -87,7 +90,7 @@ def assert_correct(actual, array): [[-1], [-1], [2], [0], [1], [0], [0]], ], ) -def test_contraction_one_argument_non_timeseries_non_voxel( +def test_contraction_one_argument( size, is_time_series, is_voxel, desired_pattern, dimension): torch_tensor = torch.rand(*size) x = phlower_tensor( @@ -303,3 +306,250 @@ def test_contraction_raises_phlower_incompatible_tensor_error(): torch.rand(10 * 10 * 10, 3, 16), is_voxel=False) with pytest.raises(PhlowerIncompatibleTensorError): _functions.contraction(x, y) + + +@pytest.mark.parametrize( + "size_x, size_y, x_is_time_series, y_is_time_series, is_voxel, " + "desired_pattern", + [ + # Base + ( + (10, 16), (10, 16), False, False, False, + "nf,nf->nf"), + ( + (10, 3, 16), (10, 16), False, False, False, + "npf,nf->npf"), + ( + (10, 16), (10, 3, 16), False, False, False, + "nf,npf->npf"), + ( + (10, 3, 16), (10, 3, 16), False, False, False, + "npf,nqf->npqf"), + ( + (10, 3, 3, 16), (10, 16), False, False, False, + "npqf,nf->npqf"), + ( + (10, 3, 3, 16), (10, 3, 16), False, False, False, + "npqf,nrf->npqrf"), + ( + (10, 3, 3, 16), (10, 3, 3, 16), False, False, False, + "npqf,nrsf->npqrsf"), + # X time series + ( + (4, 10, 16), (10, 16), True, False, False, + "tnf,nf->tnf"), + ( + (4, 10, 3, 16), (10, 16), True, False, False, + "tnpf,nf->tnpf"), + ( + (4, 10, 16), (10, 3, 16), True, False, False, + "tnf,npf->tnpf"), + ( + (4, 10, 3, 16), (10, 3, 16), True, False, False, + "tnpf,nqf->tnpqf"), + ( + (4, 10, 3, 3, 16), (10, 16), True, False, False, + "tnpqf,nf->tnpqf"), + ( + (4, 10, 3, 3, 16), (10, 3, 16), True, False, False, + "tnpqf,nrf->tnpqrf"), + ( + (4, 10, 3, 3, 16), (10, 3, 3, 16), True, False, False, + "tnpqf,nrsf->tnpqrsf"), + # Y time series + ( + (10, 16), (4, 10, 16), False, True, False, + "nf,tnf->tnf"), + ( + (10, 3, 16), (4, 10, 16), False, True, False, + "npf,tnf->tnpf"), + ( + (10, 16), (4, 10, 3, 16), False, True, False, + "nf,tnpf->tnpf"), + ( + (10, 3, 16), (4, 10, 3, 16), False, True, False, + "npf,tnqf->tnpqf"), + ( + (10, 3, 3, 16), (4, 10, 16), False, True, False, + "npqf,tnf->tnpqf"), + ( + (10, 3, 3, 16), (4, 10, 3, 16), False, True, False, + "npqf,tnrf->tnpqrf"), + ( + (10, 3, 3, 16), (4, 10, 3, 3, 16), False, True, False, + "npqf,tnrsf->tnpqrsf"), + # X Y time series + ( + (4, 10, 16), (4, 10, 16), True, True, False, + "tnf,tnf->tnf"), + ( + (4, 10, 3, 16), (4, 10, 16), True, True, False, + "tnpf,tnf->tnpf"), + ( + (4, 10, 16), (4, 10, 3, 16), True, True, False, + "tnf,tnpf->tnpf"), + ( + (4, 10, 3, 16), (4, 10, 3, 16), True, True, False, + "tnpf,tnqf->tnpqf"), + ( + (4, 10, 3, 3, 16), (4, 10, 16), True, True, False, + "tnpqf,tnf->tnpqf"), + ( + (4, 10, 3, 3, 16), (4, 10, 3, 16), True, True, False, + "tnpqf,tnrf->tnpqrf"), + ( + (4, 10, 3, 3, 16), (4, 10, 3, 3, 16), True, True, False, + "tnpqf,tnrsf->tnpqrsf"), + # X time series, X Y voxel + ( + (4, 10, 10, 10, 16), (10, 10, 10, 16), + True, False, True, "txyzf,xyzf->txyzf"), + ( + (4, 10, 10, 10, 3, 16), (10, 10, 10, 16), + True, False, True, "txyzpf,xyzf->txyzpf"), + ( + (4, 10, 10, 10, 16), (10, 10, 10, 3, 16), + True, False, True, "txyzf,xyzpf->txyzpf"), + ( + (4, 10, 10, 10, 3, 16), (10, 10, 10, 3, 16), + True, False, True, "txyzpf,xyzqf->txyzpqf"), + ( + (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 16), + True, False, True, "txyzpqf,xyzf->txyzpqf"), + ( + (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 16), + True, False, True, "txyzpqf,xyzrf->txyzpqrf"), + ( + (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 3, 16), + True, False, True, "txyzpqf,xyzrsf->txyzpqrsf"), + # X Y time series, X Y voxel + ( + (4, 10, 10, 10, 16), (4, 10, 10, 10, 16), + True, True, True, "txyzf,txyzf->txyzf"), + ( + (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 16), + True, True, True, "txyzpf,txyzf->txyzpf"), + ( + (4, 10, 10, 10, 16), (4, 10, 10, 10, 3, 16), + True, True, True, "txyzf,txyzpf->txyzpf"), + ( + (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 3, 16), + True, True, True, "txyzpf,txyzqf->txyzpqf"), + ( + (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 16), + True, True, True, "txyzpqf,txyzf->txyzpqf"), + ( + (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 16), + True, True, True, "txyzpqf,txyzrf->txyzpqrf"), + ( + (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 3, 16), + True, True, True, "txyzpqf,txyzrsf->txyzpqrsf"), + ], +) +@pytest.mark.parametrize( + "dimension_x", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +@pytest.mark.parametrize( + "dimension_y", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +def test_tensor_product( + size_x, size_y, x_is_time_series, y_is_time_series, is_voxel, + desired_pattern, dimension_x, dimension_y): + t_x = torch.rand(*size_x) + x = phlower_tensor( + t_x, dimension=dimension_x, + is_time_series=x_is_time_series, is_voxel=is_voxel) + + t_y = torch.rand(*size_y) + y = phlower_tensor( + t_y, dimension=dimension_y, + is_time_series=y_is_time_series, is_voxel=is_voxel) + + actual = _functions.tensor_product(x, y) + desired = torch.einsum( + desired_pattern, t_x, t_y).numpy() + np.testing.assert_almost_equal(actual.to_numpy(), desired) + + assert actual.is_time_series == x_is_time_series or y_is_time_series + assert actual.is_voxel == is_voxel + assert actual.rank() == x.rank() + y.rank() + + if dimension_x is not None and dimension_y is not None: + actual_dimension = actual.dimension._tensor.numpy() + desired_dimension = np.array(dimension_x) + np.array(dimension_y) + np.testing.assert_almost_equal(actual_dimension, desired_dimension) + else: + assert actual.dimension is None + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel, desired_pattern", + [ + ((10, 1), False, False, None), + ((10, 16), False, False, None), + ((10, 3, 16), False, False, "pq,nqf->npf"), + ((10, 3, 3, 16), False, False, "pq,rs,nqsf->nprf"), + ((4, 10, 1), True, False, None), + ((4, 10, 16), True, False, None), + ((4, 10, 3, 16), True, False, "pq,tnqf->tnpf"), + ((4, 10, 3, 3, 16), True, False, "pq,rs,tnqsf->tnprf"), + ((10, 10, 10, 1), False, True, None), + ((10, 10, 10, 16), False, True, None), + ((10, 10, 10, 3, 16), False, True, "pq,xyzqf->xyzpf"), + ((10, 10, 10, 3, 3, 16), False, True, "pq,rs,xyzqsf->xyzprf"), + ((4, 10, 10, 10, 1), True, True, None), + ((4, 10, 10, 10, 16), True, True, None), + ((4, 10, 10, 10, 3, 16), True, True, "pq,txyzqf->txyzpf"), + ((4, 10, 10, 10, 3, 3, 16), True, True, "pq,rs,txyzqsf->txyzprf"), + ], +) +@pytest.mark.parametrize( + "dimension", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +def test_apply_orthogonal_group( + size, is_time_series, is_voxel, desired_pattern, dimension): + orthogonal_matrix = torch.from_numpy(ortho_group.rvs(3).astype(np.float32)) + orthogonal_tensor = phlower_tensor(orthogonal_matrix) + + torch_tensor = torch.rand(*size) + x = phlower_tensor( + torch_tensor, dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + actual = _functions.apply_orthogonal_group(orthogonal_tensor, x) + + if desired_pattern is None: + desired = torch_tensor.numpy() + else: + inputs = [orthogonal_matrix] * x.rank() + [torch_tensor] + desired = torch.einsum( + desired_pattern, *inputs).numpy() + np.testing.assert_almost_equal(actual.to_numpy(), desired) + + assert actual.is_time_series == is_time_series + assert actual.is_voxel == is_voxel + assert actual.rank() == x.rank() + + if dimension is not None: + actual_dimension = actual.dimension._tensor.numpy() + desired_dimension = dimension + np.testing.assert_almost_equal(actual_dimension, desired_dimension) + else: + assert actual.dimension is None diff --git a/tests/test_nn/test_core_modules/test_identity.py b/tests/test_nn/test_core_modules/test_identity.py new file mode 100644 index 0000000..977d260 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_identity.py @@ -0,0 +1,34 @@ +import numpy as np +import pytest +import torch + +from phlower import PhlowerTensor +from phlower.collections import phlower_tensor_collection +from phlower.nn import Identity + + +def test__can_call_parameters(): + model = Identity() + + # To check Identity inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "input_shape", + [ + (5, 5, 32), + (1, 2, 48), + ], +) +def test__identity(input_shape): + phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) + phlower_tensors = phlower_tensor_collection( + {"phlower_tensor": phlower_tensor}) + + model = Identity() + + actual = model(phlower_tensors) + + np.testing.assert_almost_equal( + actual.to_numpy(), phlower_tensor.to_numpy()) diff --git a/tests/test_nn/test_core_modules/test_proportional.py b/tests/test_nn/test_core_modules/test_proportional.py new file mode 100644 index 0000000..f7cc7bd --- /dev/null +++ b/tests/test_nn/test_core_modules/test_proportional.py @@ -0,0 +1,45 @@ +import numpy as np +import pytest +import torch + +from phlower import PhlowerTensor +from phlower.collections import phlower_tensor_collection +from phlower.nn import Proportional + + +def test__can_call_parameters(): + model = Proportional(nodes=[4, 8]) + + # To check Concatenator inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "size, is_time_series", + [ + ((10, 1), False), + ((10, 16), False), + ((10, 3, 16), False), + ((4, 10, 1), True), + ((4, 10, 16), True), + ((4, 10, 3, 16), True), + ], +) +@pytest.mark.parametrize("n_output_feature", [1, 16, 32]) +@pytest.mark.parametrize("scale", [0., 0.5, 2.]) +def test__proportional_linearity( + size, is_time_series, n_output_feature, scale): + + model = Proportional(nodes=[size[-1], n_output_feature]) + + phlower_tensor = PhlowerTensor( + torch.rand(*size), is_time_series=is_time_series) + + phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) + actual = model(phlower_tensors).to_numpy() + + scaled_phlower_tensors = phlower_tensor_collection( + {'tensor': phlower_tensor * scale}) + desired = model(scaled_phlower_tensors).to_numpy() + + np.testing.assert_almost_equal(actual * scale, desired) From 878108e9ead14759325db55ce9e23fbe5645d6c7 Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 05:46:13 +0900 Subject: [PATCH 24/89] add PInvMLP --- src/phlower/_base/tensors/_phlower_tensor.py | 25 +++ src/phlower/nn/__init__.py | 1 + src/phlower/nn/_core_modules/__init__.py | 2 + src/phlower/nn/_core_modules/_functions.py | 17 ++ src/phlower/nn/_core_modules/_pinv_mlp.py | 154 ++++++++++++++++++ src/phlower/nn/_core_modules/_utils.py | 3 + .../settings/_module_settings/__init__.py | 2 + .../_module_settings/_pinv_mlp_setting.py | 50 ++++++ src/phlower/utils/exceptions.py | 3 + .../test_core_modules/test_functions.py | 12 ++ tests/test_nn/test_core_modules/test_pinv.py | 42 +++++ 11 files changed, 311 insertions(+) create mode 100644 src/phlower/nn/_core_modules/_pinv_mlp.py create mode 100644 src/phlower/settings/_module_settings/_pinv_mlp_setting.py create mode 100644 tests/test_nn/test_core_modules/test_pinv.py diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 84e1a07..f1c7367 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -139,18 +139,43 @@ def __str__(self) -> str: f"Dimension: {self._dimension_tensor})" ) + def __eq__(self, other): + return torch.eq(self, other) + + def __lt__(self, other): + return torch.lt(self, other) + + def __le__(self, other): + return torch.le(self, other) + + def __gt__(self, other): + return torch.gt(self, other) + + def __ge__(self, other): + return torch.ge(self, other) + def __abs__(self) -> PhlowerTensor: return torch.abs(self) def __sub__(self, other: PhlowerTensor): return torch.subtract(self, other) + def __neg__(self): + return torch.neg(self) + def __add__(self, other) -> PhlowerTensor: return torch.add(self, other) def __mul__(self, other) -> PhlowerTensor: return torch.mul(self, other) + def __setitem__(self, key, value): + if isinstance(key, PhlowerTensor): + self._tensor[key.to_tensor()] = value + else: + self._tensor[key] = value + return self + def __len__(self) -> int: return len(self._tensor) diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 17f1131..7518e1c 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -2,6 +2,7 @@ from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._identity import Identity from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._pinv_mlp import PInvMLP from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share from phlower.nn._group_module import PhlowerGroupModule diff --git a/src/phlower/nn/_core_modules/__init__.py b/src/phlower/nn/_core_modules/__init__.py index 65ecce9..544092a 100644 --- a/src/phlower/nn/_core_modules/__init__.py +++ b/src/phlower/nn/_core_modules/__init__.py @@ -2,6 +2,7 @@ from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._identity import Identity from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._pinv_mlp import PInvMLP from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share from phlower.nn._interface_module import IPhlowerCoreModule @@ -16,6 +17,7 @@ GCN, Identity, MLP, + PInvMLP, Proportional, Share, ] diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 7088c37..3df6258 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -12,6 +12,23 @@ def identity(x): return x +def leaky_relu0p5(x): + """Leaky ReLU with the negative slope = 0.5.""" + return torch.nn.functional.leaky_relu(x, negative_slope=0.5) + + +def inversed_leaky_relu0p5(x): + """Inverse of leaky_relu0p5.""" + return torch.nn.functional.leaky_relu(x, negative_slope=2) + + +def truncated_atanh(x, epsilon=1e-8): + """Inverse tanh with truncating values >=1 or <=-1.""" + x[x>=1. - epsilon] = 1. - epsilon + x[x<=-1. + epsilon] = -1. + epsilon + return torch.atanh(x) + + def spmm( sparse: IPhlowerTensor, x: IPhlowerTensor, repeat: int = 1) -> IPhlowerTensor: diff --git a/src/phlower/nn/_core_modules/_pinv_mlp.py b/src/phlower/nn/_core_modules/_pinv_mlp.py new file mode 100644 index 0000000..15ca2be --- /dev/null +++ b/src/phlower/nn/_core_modules/_pinv_mlp.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import torch +from typing_extensions import Self + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._core_modules import _utils +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings import PInvMLPSetting +from phlower.utils.exceptions import ( + NotFoundReferenceModuleError, + PhlowerInvalidActivationError, +) + + +class PInvMLP(IPhlowerCoreModule, torch.nn.Module): + """ + Pseudo inverse of the reference MLP layer. + """ + + @classmethod + def from_setting(cls, setting: PInvMLPSetting) -> Self: + return cls(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + return "PInvMLP" + + @classmethod + def need_reference(cls) -> bool: + return True + + def __init__(self, reference_name: str, **kwards) -> None: + super().__init__() + + self._reference_name = reference_name + self._reference: IPhlowerCoreModule | None = None + + def get_reference_name(self) -> str: + return self._reference_name + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: + assert parent is not None + + try: + self._reference = parent.search_module(self._reference_name) + except KeyError as ex: + raise NotFoundReferenceModuleError( + f"Reference module {self._reference_name} is not found " + "in the same group." + ) from ex + + self._initialize() + + def _initialize(self): + """Initialize parameters and activations after resolve() is called.""" + self._activation_names = [ + self._inverse_activation_name(a) + for a in self._reference._activations[::-1]] + self._activations = [ + _utils.ActivationSelector.select(name) + for name in self._activation_names] + self._chains = self._init_pinv_chains() + + def _inverse_activation_name(self, activation_name): + if activation_name == "identity": + return "identity" + if activation_name == "leaky_relu0p5": + return "inversed_leaky_relu0p5" + if activation_name == "tanh": + return "truncated_atanh" + + raise PhlowerInvalidActivationError( + f"Cannot pinv for {activation_name}") + + def _init_pinv_chains(self): + name = self._reference.__class__.__name__ + if name in ["MLP", "Proportional"]: + return self._init_pinv_mlp_chains(self._reference._chains) + + raise ValueError(f"Unsupported reference class: {name}") + + def _init_pinv_mlp_chains( + self, chains: _utils.ExtendedLinearList, option=None): + return [PInvLinear(c, option=option) for c in chains._linears[::-1]] + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + **kwards, + ) -> PhlowerTensor: + """forward function which overload torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor]): + sparse tensor objects + + Returns: + PhlowerTensor: + Tensor object + """ + if self._reference is None: + raise ValueError( + "reference module in PInvMLP module is not set. " + "Please check that `resolve` function is called." + ) + + h = data.unique_item() + for activation, chain in zip( + self._activations, self._chains, strict=True): + # Activation comes first because it is pseudo inverse + h = chain(activation(h)) + + return h + + +class PInvLinear(torch.nn.Module): + + def __init__(self, ref_linear: torch.nn.Linear, option: str | None = None): + super().__init__() + self.ref_linear = ref_linear + self.option = option + return + + def forward(self, x): + h = torch.nn.functional.linear(x + self.bias, self.weight) + return h + + @property + def weight(self): + """Return pseudo inversed weight.""" + if self.option is None: + w = self.ref_linear.weight + else: + raise ValueError(f"Unexpected option: {self.option}") + return torch.pinverse(w) + + @property + def bias(self): + """Return inverse bias.""" + if self.ref_linear.bias is None: + return 0 + else: + return - self.ref.bias diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index 2a038d4..ed59743 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -77,10 +77,13 @@ def _validate_args(self) -> None: class ActivationSelector: _REGISTERED_ACTIVATIONS = { "identity": _functions.identity, + "inversed_leaky_relu0p5": _functions.inversed_leaky_relu0p5, + "leaky_relu0p5": _functions.leaky_relu0p5, "relu": torch.relu, "sigmoid": torch.sigmoid, "sqrt": torch.sqrt, "tanh": torch.tanh, + "truncated_atanh": _functions.truncated_atanh, } @staticmethod diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index 73e351f..f0dfd2d 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -11,6 +11,7 @@ from phlower.settings._module_settings._gcn_setting import GCNSetting from phlower.settings._module_settings._identity_setting import IdentitySetting from phlower.settings._module_settings._mlp_setting import MLPSetting +from phlower.settings._module_settings._pinv_mlp_setting import PInvMLPSetting from phlower.settings._module_settings._proportional_setting import ( ProportionalSetting, ) @@ -22,6 +23,7 @@ "GCN": GCNSetting, "Identity": IdentitySetting, "MLP": MLPSetting, + "PInvMLP": PInvMLPSetting, "Proportional": ProportionalSetting, "Share": ShareSetting, } diff --git a/src/phlower/settings/_module_settings/_pinv_mlp_setting.py b/src/phlower/settings/_module_settings/_pinv_mlp_setting.py new file mode 100644 index 0000000..bce9cfc --- /dev/null +++ b/src/phlower/settings/_module_settings/_pinv_mlp_setting.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import pydantic +from pydantic import Field + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) +from phlower.utils.exceptions import NotFoundReferenceModuleError + + +class PInvMLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): + reference_name: str = Field(..., frozen=True) + reference: IPhlowerLayerParameters | None = Field(None, exclude=True) + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict( + extra="forbid", arbitrary_types_allowed=True + ) + + def gather_input_dims(self, *input_dims: int) -> int: + return self.reference.gather_input_dims(*input_dims) + + def get_n_nodes(self) -> list[int] | None: + return self.reference.get_n_nodes()[::-1] + + def check_exist_reference(self) -> None: + if self.reference is not None: + return + + raise ValueError( + f"Reference setting {self.reference_name} in PinvMLP is None." + "Please check that `get_reference` method has been called." + ) + + def overwrite_nodes(self, nodes: list[int]) -> None: ... + + @property + def need_reference(self) -> bool: + return True + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + try: + self.reference = parent.search_module_setting(self.reference_name) + except KeyError as ex: + raise NotFoundReferenceModuleError( + f"Reference module {self.reference_name} is not found " + "in the same group." + ) from ex diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 11ea81b..70a78b7 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -74,3 +74,6 @@ class PhlowerIncompatibleTensorError(ValueError): class PhlowerTypeError(TypeError): """This error raises when type is not compatible with what is expected.""" + +class PhlowerInvalidActivationError(ValueError): + """This error raises when a set activation is invalid""" diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index cb6c062..a4a0195 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -60,6 +60,18 @@ def assert_correct(actual, array): assert_correct(actual_spmm, np_dense) +def test_leaky_relu0p5_inverse_leaky_relu0p5(): + x = PhlowerTensor(torch.rand(100)) + y = _functions.inversed_leaky_relu0p5(_functions.leaky_relu0p5(x)) + np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) + + +def test_tanh_truncated_atanh(): + x = PhlowerTensor(torch.rand(100)) + y = _functions.truncated_atanh(torch.tanh(x)) + np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) + + @pytest.mark.parametrize( "size, is_time_series, is_voxel, desired_pattern", [ diff --git a/tests/test_nn/test_core_modules/test_pinv.py b/tests/test_nn/test_core_modules/test_pinv.py new file mode 100644 index 0000000..1f33d4d --- /dev/null +++ b/tests/test_nn/test_core_modules/test_pinv.py @@ -0,0 +1,42 @@ + +import numpy as np +import pytest +import torch + +from phlower import PhlowerTensor +from phlower.collections import phlower_tensor_collection +from phlower.nn import MLP, PinvMLP + + +def test__can_call_parameters(): + model = PinvMLP(reference_name="MLP0") + MLP0 = MLP(nodes=[10, 10]) + model._reference = MLP0 + + # To check Concatenator inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "mlp_nodes, activations", + [ + ([10, 10], ["identity"]), + ([10, 12], ["leaky_relu0p5"]), + ([20, 40, 100], ["tanh", "identity"]) + ], +) +def test__pinv_mlp(mlp_nodes, activations): + MLP0 = MLP(nodes=mlp_nodes, activations=activations) + + model = PinvMLP(reference_name="MLP0") + model._reference = MLP0 + model._initialize() + + t = PhlowerTensor(tensor=torch.rand(10, 3, mlp_nodes[0])) + phlower_tensors = phlower_tensor_collection({"tensor": t}) + + mlp_val = MLP0(phlower_tensors) + pinv_val = model(phlower_tensor_collection({"tensor": mlp_val})) + + np.testing.assert_array_almost_equal( + pinv_val.to_numpy(), t.to_numpy(), decimal=5) From 3b875f0a58f143a3837870fc773b549d8034df76 Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 06:35:37 +0900 Subject: [PATCH 25/89] add more activations --- src/phlower/_base/tensors/_phlower_tensor.py | 6 ++++++ src/phlower/nn/_core_modules/_functions.py | 16 ++++++++++++++++ src/phlower/nn/_core_modules/_pinv_mlp.py | 2 ++ src/phlower/nn/_core_modules/_utils.py | 2 ++ .../test_core_modules/test_functions.py | 10 ++++++++-- tests/test_nn/test_core_modules/test_pinv.py | 19 ++++++++++--------- 6 files changed, 44 insertions(+), 11 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index f1c7367..7f3172e 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -169,6 +169,12 @@ def __add__(self, other) -> PhlowerTensor: def __mul__(self, other) -> PhlowerTensor: return torch.mul(self, other) + def __rmul__(self, other) -> PhlowerTensor: + return torch.mul(self, other) + + def __pow__(self, other) -> PhlowerTensor: + return torch.pow(self, other) + def __setitem__(self, key, value): if isinstance(key, PhlowerTensor): self._tensor[key.to_tensor()] = value diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 3df6258..f1801f5 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -29,6 +29,22 @@ def truncated_atanh(x, epsilon=1e-8): return torch.atanh(x) +def smooth_leaky_relu(x): + """Smooth leaky ReLU""" + # a x + (1 - a) x sqrt(x**2 + b) + # return 0.75 * x + 0.25 * (x**2 + 1 / 16)**.5 # b = 1 / 16 + return 0.75 * x + 0.25 * (x**2 + 1 / 100)**.5 # b = 1 / 100 + + +def inversed_smooth_leaky_relu(x): + # return 1.5 * x - 1 / 16 * torch.sqrt((8 * x)**2 + 2) # b = 1 / 16 + return 1.5 * x - 1 / 40 * torch.sqrt(400 * x**2 + 2) # b = 1 / 100 + + +def derivative_smooth_leaky_relu(x): + # return 0.75 + 0.25 * x / torch.sqrt(x**2 + 1 / 16) # b = 1 / 16 + return 0.75 + 0.25 * x / torch.sqrt(x**2 + 1 / 100) + def spmm( sparse: IPhlowerTensor, x: IPhlowerTensor, repeat: int = 1) -> IPhlowerTensor: diff --git a/src/phlower/nn/_core_modules/_pinv_mlp.py b/src/phlower/nn/_core_modules/_pinv_mlp.py index 15ca2be..d474912 100644 --- a/src/phlower/nn/_core_modules/_pinv_mlp.py +++ b/src/phlower/nn/_core_modules/_pinv_mlp.py @@ -73,6 +73,8 @@ def _inverse_activation_name(self, activation_name): return "identity" if activation_name == "leaky_relu0p5": return "inversed_leaky_relu0p5" + if activation_name == "smooth_leaky_relu": + return "inversed_smooth_leaky_relu" if activation_name == "tanh": return "truncated_atanh" diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index ed59743..a2407ed 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -78,9 +78,11 @@ class ActivationSelector: _REGISTERED_ACTIVATIONS = { "identity": _functions.identity, "inversed_leaky_relu0p5": _functions.inversed_leaky_relu0p5, + "inversed_smooth_leaky_relu": _functions.inversed_smooth_leaky_relu, "leaky_relu0p5": _functions.leaky_relu0p5, "relu": torch.relu, "sigmoid": torch.sigmoid, + "smooth_leaky_relu": _functions.smooth_leaky_relu, "sqrt": torch.sqrt, "tanh": torch.tanh, "truncated_atanh": _functions.truncated_atanh, diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index a4a0195..c5cf50f 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -61,17 +61,23 @@ def assert_correct(actual, array): def test_leaky_relu0p5_inverse_leaky_relu0p5(): - x = PhlowerTensor(torch.rand(100)) + x = PhlowerTensor(torch.rand(100)) * 4 - 2. y = _functions.inversed_leaky_relu0p5(_functions.leaky_relu0p5(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) def test_tanh_truncated_atanh(): - x = PhlowerTensor(torch.rand(100)) + x = PhlowerTensor(torch.rand(100)) * 4 - 2 y = _functions.truncated_atanh(torch.tanh(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) +def test_smooth_leaky_relu_inverse(): + x = PhlowerTensor(torch.rand(100)) * 4 - 2 + y = _functions.inversed_smooth_leaky_relu(_functions.smooth_leaky_relu(x)) + np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy(), decimal=5) + + @pytest.mark.parametrize( "size, is_time_series, is_voxel, desired_pattern", [ diff --git a/tests/test_nn/test_core_modules/test_pinv.py b/tests/test_nn/test_core_modules/test_pinv.py index 1f33d4d..1697f03 100644 --- a/tests/test_nn/test_core_modules/test_pinv.py +++ b/tests/test_nn/test_core_modules/test_pinv.py @@ -5,11 +5,11 @@ from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection -from phlower.nn import MLP, PinvMLP +from phlower.nn import MLP, PInvMLP def test__can_call_parameters(): - model = PinvMLP(reference_name="MLP0") + model = PInvMLP(reference_name="MLP0") MLP0 = MLP(nodes=[10, 10]) model._reference = MLP0 @@ -18,17 +18,18 @@ def test__can_call_parameters(): @pytest.mark.parametrize( - "mlp_nodes, activations", + "mlp_nodes, activations, decimal", [ - ([10, 10], ["identity"]), - ([10, 12], ["leaky_relu0p5"]), - ([20, 40, 100], ["tanh", "identity"]) + ([10, 10], ["identity"], 5), + ([10, 12], ["leaky_relu0p5"], 5), + ([20, 40, 100], ["tanh", "identity"], 5), + ([20, 20, 40, 100], ["tanh", "smooth_leaky_relu", "leaky_relu0p5"], 3), ], ) -def test__pinv_mlp(mlp_nodes, activations): +def test__pinv_mlp(mlp_nodes, activations, decimal): MLP0 = MLP(nodes=mlp_nodes, activations=activations) - model = PinvMLP(reference_name="MLP0") + model = PInvMLP(reference_name="MLP0") model._reference = MLP0 model._initialize() @@ -39,4 +40,4 @@ def test__pinv_mlp(mlp_nodes, activations): pinv_val = model(phlower_tensor_collection({"tensor": mlp_val})) np.testing.assert_array_almost_equal( - pinv_val.to_numpy(), t.to_numpy(), decimal=5) + pinv_val.to_numpy(), t.to_numpy(), decimal=decimal) From 5fa484d43cdc69df86c121a3ca150e528c98f870 Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 06:36:50 +0900 Subject: [PATCH 26/89] add more comment --- src/phlower/nn/_core_modules/_functions.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index f1801f5..85e88c3 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -32,11 +32,13 @@ def truncated_atanh(x, epsilon=1e-8): def smooth_leaky_relu(x): """Smooth leaky ReLU""" # a x + (1 - a) x sqrt(x**2 + b) + # return 0.75 * x + 0.25 * (x**2 + 8)**.5 # b = 8 # return 0.75 * x + 0.25 * (x**2 + 1 / 16)**.5 # b = 1 / 16 return 0.75 * x + 0.25 * (x**2 + 1 / 100)**.5 # b = 1 / 100 def inversed_smooth_leaky_relu(x): + # return 1.5 * x - ((0.5 * x)**2 + 1)**.5 # b = 8 # return 1.5 * x - 1 / 16 * torch.sqrt((8 * x)**2 + 2) # b = 1 / 16 return 1.5 * x - 1 / 40 * torch.sqrt(400 * x**2 + 2) # b = 1 / 100 From d830b388310bd5a9468c7fd19dbfcc7944c71373 Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 06:49:33 +0900 Subject: [PATCH 27/89] update test --- tests/test_nn/test_core_modules/test_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index c5cf50f..3d283a9 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -69,7 +69,7 @@ def test_leaky_relu0p5_inverse_leaky_relu0p5(): def test_tanh_truncated_atanh(): x = PhlowerTensor(torch.rand(100)) * 4 - 2 y = _functions.truncated_atanh(torch.tanh(x)) - np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) + np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy(), decimal=6) def test_smooth_leaky_relu_inverse(): From ddef85f3994becae7d5de513c5d8ca8a0f411f59 Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 17:39:09 +0900 Subject: [PATCH 28/89] update smooth leaky relu --- src/phlower/_base/tensors/_phlower_tensor.py | 6 ++++ src/phlower/nn/_core_modules/_functions.py | 28 +++++++++++-------- src/phlower/nn/_core_modules/_utils.py | 6 ++-- .../test_core_modules/test_functions.py | 3 +- 4 files changed, 28 insertions(+), 15 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 7f3172e..5f5efa7 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -166,12 +166,18 @@ def __neg__(self): def __add__(self, other) -> PhlowerTensor: return torch.add(self, other) + def __radd__(self, other) -> PhlowerTensor: + return torch.add(self, other) + def __mul__(self, other) -> PhlowerTensor: return torch.mul(self, other) def __rmul__(self, other) -> PhlowerTensor: return torch.mul(self, other) + def __truediv__(self, other) -> PhlowerTensor: + return torch.div(self, other) + def __pow__(self, other) -> PhlowerTensor: return torch.pow(self, other) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 85e88c3..07c8dae 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -29,23 +29,27 @@ def truncated_atanh(x, epsilon=1e-8): return torch.atanh(x) -def smooth_leaky_relu(x): +class SmoothLeakyReLU(): """Smooth leaky ReLU""" - # a x + (1 - a) x sqrt(x**2 + b) - # return 0.75 * x + 0.25 * (x**2 + 8)**.5 # b = 8 - # return 0.75 * x + 0.25 * (x**2 + 1 / 16)**.5 # b = 1 / 16 - return 0.75 * x + 0.25 * (x**2 + 1 / 100)**.5 # b = 1 / 100 + def __init__(self, a: float = 0.75, b: float = 1 / 100): + self.a = a + self.b = b -def inversed_smooth_leaky_relu(x): - # return 1.5 * x - ((0.5 * x)**2 + 1)**.5 # b = 8 - # return 1.5 * x - 1 / 16 * torch.sqrt((8 * x)**2 + 2) # b = 1 / 16 - return 1.5 * x - 1 / 40 * torch.sqrt(400 * x**2 + 2) # b = 1 / 100 + def __call__(self, x): + # a x + (1 - a) sqrt(x**2 + b) + return self.a * x + (1 - self.a) * torch.sqrt(x**2 + self.b) + def inverse(self, x): + return ( + self.a * x + - torch.sqrt( + (self.a - 1)**2 * (2 * self.a * self.b - self.b + x**2)) + ) / (2 * self.a - 1) + + def derivative(self, x): + return self.a - ((self.a - 1) * x) / torch.sqrt(self.b + x**2) -def derivative_smooth_leaky_relu(x): - # return 0.75 + 0.25 * x / torch.sqrt(x**2 + 1 / 16) # b = 1 / 16 - return 0.75 + 0.25 * x / torch.sqrt(x**2 + 1 / 100) def spmm( sparse: IPhlowerTensor, x: IPhlowerTensor, diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index a2407ed..e625074 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -75,14 +75,16 @@ def _validate_args(self) -> None: class ActivationSelector: + _SMOOTH_LEAKY_RELU = _functions.SmoothLeakyReLU() + _REGISTERED_ACTIVATIONS = { "identity": _functions.identity, "inversed_leaky_relu0p5": _functions.inversed_leaky_relu0p5, - "inversed_smooth_leaky_relu": _functions.inversed_smooth_leaky_relu, + "inversed_smooth_leaky_relu": _SMOOTH_LEAKY_RELU.inverse, "leaky_relu0p5": _functions.leaky_relu0p5, "relu": torch.relu, "sigmoid": torch.sigmoid, - "smooth_leaky_relu": _functions.smooth_leaky_relu, + "smooth_leaky_relu": _SMOOTH_LEAKY_RELU, "sqrt": torch.sqrt, "tanh": torch.tanh, "truncated_atanh": _functions.truncated_atanh, diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 3d283a9..892ac0c 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -74,7 +74,8 @@ def test_tanh_truncated_atanh(): def test_smooth_leaky_relu_inverse(): x = PhlowerTensor(torch.rand(100)) * 4 - 2 - y = _functions.inversed_smooth_leaky_relu(_functions.smooth_leaky_relu(x)) + f = _functions.SmoothLeakyReLU() + y = f.inverse(f(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy(), decimal=5) From 5bd2fc14c34adc939b109093150b6ae77977e27f Mon Sep 17 00:00:00 2001 From: horiem Date: Sun, 11 Aug 2024 17:56:47 +0900 Subject: [PATCH 29/89] fix for lint --- src/phlower/nn/_core_modules/_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 07c8dae..d4c9bcf 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -29,7 +29,7 @@ def truncated_atanh(x, epsilon=1e-8): return torch.atanh(x) -class SmoothLeakyReLU(): +class SmoothLeakyReLU: """Smooth leaky ReLU""" def __init__(self, a: float = 0.75, b: float = 1 / 100): From 9fc9be722ebf905116cb7cb5820eae2c1dce9f87 Mon Sep 17 00:00:00 2001 From: horiem Date: Mon, 12 Aug 2024 03:02:37 +0900 Subject: [PATCH 30/89] add SimilarityEquivariantMLP --- .../_base/tensors/_dimension_tensor.py | 57 +++++- src/phlower/_base/tensors/_phlower_tensor.py | 8 +- .../tensors/_tensor_collections.py | 9 + src/phlower/nn/__init__.py | 7 +- src/phlower/nn/_core_modules/__init__.py | 3 + src/phlower/nn/_core_modules/_functions.py | 55 ++++++ .../_similarity_equivariant_mlp.py | 178 ++++++++++++++++++ .../settings/_module_settings/__init__.py | 3 + .../_similarity_equivariant_mlp_setting.py | 78 ++++++++ src/phlower/utils/exceptions.py | 8 + .../test_tensors/test__phlower_tensor.py | 60 ++++++ .../test_core_modules/test_functions.py | 54 ++++++ tests/test_nn/test_core_modules/test_pinv.py | 6 +- .../test_similarity_equivariant_mlp.py | 155 +++++++++++++++ 14 files changed, 671 insertions(+), 10 deletions(-) create mode 100644 src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py create mode 100644 src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py create mode 100644 tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index b0eddb6..08488ab 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -113,6 +113,10 @@ def to(self, device: str | torch.device, non_blocking: bool = False): def detach(self) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(tensor=self._tensor.detach()) + def is_dimensionless(self) -> bool: + """Return True if the tensor is dimensionless.""" + return torch.sum(torch.abs(self._tensor)) < 1e-5 + @classmethod def __torch_function__(cls, func, types, args: tuple, kwargs=None): if kwargs is None: @@ -145,8 +149,21 @@ def add(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): if inputs != other: raise DimensionIncompatibleError( - "Add operation for different physical dimensions is not allowed." - ) + "Add operation for different physical dimensions is not " + "allowed.") + + return PhlowerDimensionTensor(inputs._tensor) + + raise DimensionIncompatibleError() + + +@dimension_wrap_implements(torch.sub) +def sub(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): + if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): + if inputs != other: + raise DimensionIncompatibleError( + "Sub operation for different physical dimensions is not " + "allowed.") return PhlowerDimensionTensor(inputs._tensor) @@ -158,6 +175,11 @@ def pow(inputs: PhlowerDimensionTensor, other): return PhlowerDimensionTensor(inputs._tensor * other) +@dimension_wrap_implements(torch.sqrt) +def sqrt(inputs: PhlowerDimensionTensor): + return PhlowerDimensionTensor(inputs._tensor / 2) + + @dimension_wrap_implements(torch.mul) def mul(inputs, other): _input = ( @@ -173,6 +195,21 @@ def mul(inputs, other): return PhlowerDimensionTensor(_input._tensor + _other._tensor) +@dimension_wrap_implements(torch.div) +def div(inputs, other): + _input = ( + inputs + if isinstance(inputs, PhlowerDimensionTensor) + else zero_dimension_tensor() + ) + _other = ( + other + if isinstance(other, PhlowerDimensionTensor) + else zero_dimension_tensor() + ) + return PhlowerDimensionTensor(_input._tensor - _other._tensor) + + @dimension_wrap_implements(torch.reshape) def reshape(inputs, shape): return PhlowerDimensionTensor(inputs._tensor) @@ -251,3 +288,19 @@ def concatenate(inputs, *args, **kwards): return PhlowerDimensionTensor(inputs[0]._tensor) raise DimensionIncompatibleError() + + +@dimension_wrap_implements(torch.tanh) +def tanh(tensor: PhlowerDimensionTensor): + if not tensor.is_dimensionless: + raise DimensionIncompatibleError( + f"Should be dimensionless to apply tanh but {tensor}") + return tensor + + +@dimension_wrap_implements(torch.nn.functional.leaky_relu) +def leaky_relu(tensor: PhlowerDimensionTensor, *args, **kwargs): + if not tensor.is_dimensionless: + raise DimensionIncompatibleError( + f"Should be dimensionless to apply leaky_relu but {tensor}") + return tensor diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 5f5efa7..a000c65 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -158,7 +158,7 @@ def __abs__(self) -> PhlowerTensor: return torch.abs(self) def __sub__(self, other: PhlowerTensor): - return torch.subtract(self, other) + return torch.sub(self, other) def __neg__(self): return torch.neg(self) @@ -178,6 +178,9 @@ def __rmul__(self, other) -> PhlowerTensor: def __truediv__(self, other) -> PhlowerTensor: return torch.div(self, other) + def __rtruediv__(self, other) -> PhlowerTensor: + return torch.div(other, self) + def __pow__(self, other) -> PhlowerTensor: return torch.pow(self, other) @@ -203,6 +206,9 @@ def coalesce(self) -> torch.Tensor: def size(self) -> torch.Size: return self._tensor.size() + def numel(self) -> int: + return torch.numel(self._tensor) + def rank(self) -> int: """Returns the tensor rank.""" if self.is_sparse: diff --git a/src/phlower/collections/tensors/_tensor_collections.py b/src/phlower/collections/tensors/_tensor_collections.py index 3232a8b..473ac76 100644 --- a/src/phlower/collections/tensors/_tensor_collections.py +++ b/src/phlower/collections/tensors/_tensor_collections.py @@ -46,6 +46,9 @@ def keys(self) -> Iterable[str]: ... @abc.abstractmethod def values(self): ... + @abc.abstractmethod + def pop(self): ... + @abc.abstractmethod def sum(self, weights: dict[str, float] = None) -> PhlowerTensor: ... @@ -92,6 +95,12 @@ def values(self): def keys(self) -> Iterable[str]: return self._x.keys() + def items(self) -> abc.ItemsView: + return self._x.items() + + def pop(self, key: str, default: PhlowerTensor = None) -> PhlowerTensor: + return self._x.pop(key, default) + def __getitem__(self, key: Any) -> PhlowerTensor: if isinstance(key, str): return self._x[key] diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 7518e1c..5f4d31f 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -1,13 +1,12 @@ from phlower.nn._core_modules._concatenator import Concatenator +from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._identity import Identity from phlower.nn._core_modules._mlp import MLP from phlower.nn._core_modules._pinv_mlp import PInvMLP from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share +from phlower.nn._core_modules \ + ._similarity_equivariant_mlp import SimilarityEquivariantMLP from phlower.nn._group_module import PhlowerGroupModule from phlower.nn._interface_module import IPhlowerCoreModule - -if True: - # NOTE: Import advanced models after - from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP diff --git a/src/phlower/nn/_core_modules/__init__.py b/src/phlower/nn/_core_modules/__init__.py index 544092a..5774589 100644 --- a/src/phlower/nn/_core_modules/__init__.py +++ b/src/phlower/nn/_core_modules/__init__.py @@ -10,6 +10,8 @@ if True: # NOTE: Import advanced models after from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP + from phlower.nn._core_modules \ + ._similarity_equivariant_mlp import SimilarityEquivariantMLP _all_models: list[IPhlowerCoreModule] = [ Concatenator, @@ -20,6 +22,7 @@ PInvMLP, Proportional, Share, + SimilarityEquivariantMLP, ] _name2model = {cls.get_nn_name(): cls for cls in _all_models} diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index d4c9bcf..3f9d09b 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -250,6 +250,29 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: is_time_series=ret_is_time_series, is_voxel=is_voxel) +def tensor_times_scalar( + tensor: IPhlowerTensor, scalar: int | float | IPhlowerTensor): + """ + Compute multiplication between tensor and scalar (field). + + Args: + tensor: IPhlowerTensor + Input tensor. + scalar: int | float | IPhlowerTensor + Input scalar or scalar field. + + Returns: + IPhlowerTensor: + Resultant tensor. + """ + if isinstance(scalar, int | float): + return tensor * scalar + if scalar.numel() == 1: + return tensor * scalar + + return tensor_product(tensor, scalar) + + def apply_orthogonal_group( orthogonal_matrix: IPhlowerTensor, tensor: IPhlowerTensor ) -> IPhlowerTensor: @@ -293,3 +316,35 @@ def apply_orthogonal_group( equation, *args, dimension=tensor.dimension, is_time_series=tensor.is_time_series, is_voxel=tensor.is_voxel) + + +def spatial_sum(tensor: IPhlowerTensor) -> IPhlowerTensor: + """Compute sum over space.""" + + if tensor.is_time_series: + time = "t" + start_space = 1 + else: + time = "" + start_space = 0 + if tensor.is_voxel: + space = "xyz" + space_width = 3 + else: + space = "x" + space_width = 1 + + squeezed = einsum( + f"{time}{space}...->{time}...", tensor, + dimension=tensor.dimension, is_time_series=tensor.is_time_series, + is_voxel=tensor.is_voxel) + + # keepdim + for _ in range(space_width): + squeezed._tensor = torch.unsqueeze(squeezed._tensor, start_space) + return squeezed + + +def spatial_mean(tensor: IPhlowerTensor) -> IPhlowerTensor: + """Compute mean over space.""" + return spatial_sum(tensor) / tensor.n_vertices() diff --git a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py new file mode 100644 index 0000000..9c035f3 --- /dev/null +++ b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import numpy as np +import torch +from typing_extensions import Self + +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import ( + IPhlowerTensorCollections, + phlower_tensor_collection, +) +from phlower.nn._core_modules import ( + MLP, + EnEquivariantMLP, + Identity, + Proportional, + _functions, +) +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings import SimilarityEquivariantMLPSetting +from phlower.utils.enums import PhysicalDimensionSymbolType +from phlower.utils.exceptions import ( + PhlowerDimensionRequiredError, + PhlowerInvalidArgumentsError, +) + + +class SimilarityEquivariantMLP(IPhlowerCoreModule, torch.nn.Module): + """ + Similarity-equivariant Multi Layer Perceptron as in + https://proceedings.mlr.press/v235/horie24a.html. + """ + + @classmethod + def from_setting(cls, setting: SimilarityEquivariantMLPSetting) -> Self: + """Generate model from setting object + + Args: + setting (SimilarityEquivariantMLPSetting): setting object + + Returns: + Self: SimilarityEquivariantMLP object + """ + return cls(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + """Return neural network name + + Returns: + str: name + """ + return "SimilarityEquivariantMLP" + + @classmethod + def need_reference(cls) -> bool: + return False + + def __init__( + self, + nodes: list[int], + activations: list[str] | None = None, + dropouts: list[float] | None = None, + bias: bool = False, + create_linear_weight: bool = False, + norm_function_name: str = None, + disable_en_equivariance: bool = False, + invariant: bool = False, + centering: bool = False, + ) -> None: + super().__init__() + + self._nodes = nodes + self._disable_en_equivariance = disable_en_equivariance + self._invariant = invariant + self._centering = centering + self._create_linear_weight = create_linear_weight + + if self._disable_en_equivariance: + self._mlp = MLP( + nodes=nodes, activations=activations, + dropouts=dropouts, bias=bias, + ) + self._linear_weight = self._init_linear_weight() + else: + self._mlp = EnEquivariantMLP( + nodes=nodes, activations=activations, + dropouts=dropouts, bias=bias, + create_linear_weight=create_linear_weight, + norm_function_name=norm_function_name, + ) + self._linear_weight = self._mlp._linear_weight + + def _init_linear_weight(self): + if not self._create_linear_weight: + if self._nodes[0] != self._nodes[-1]: + raise ValueError( + "First and last nodes are different. " + "Set create_linear_weight True.") + return Identity() + + return Proportional([self._nodes[0], self._nodes[-1]]) + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + + def get_reference_name(self) -> str | None: + return None + + def forward( + self, + data: IPhlowerTensorCollections, + *, + supports: dict[str, PhlowerTensor] | None = None, + **kwards, + ) -> PhlowerTensor: + """forward function which overloads torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + Data with the keys in tensor or PhysicalDimensionSymbolType, + with tensor is the input tensor to be computed. + PhysicalDimensionSymbolType keys denote physical scales. + No need to input all dimensions, but at least one scale + information should be input. + supports (dict[str, PhlowerTensor], optional): + Graph object. Defaults to None. + + Returns: + PhlowerTensor: Tensor object + """ + # TODO: Check if we can assume key convention + h = data.pop("tensor") + if not h.has_dimension: + raise PhlowerDimensionRequiredError("Dimension is required") + + if len(data) < 1: + raise PhlowerInvalidArgumentsError("Scale inputs are required") + if not np.all([ + PhysicalDimensionSymbolType.is_exist(k) for k in data.keys()]): + raise PhlowerInvalidArgumentsError( + "keys should be in PhysicalDimensionSymbolType. " + f"Given: {data.keys()}") + dict_dimension = h.dimension.to_dict() + + dict_scales = { + k: v**dict_dimension[k] for k, v in data.items()} + + # Make h dimensionless + for v in dict_scales.values(): + h = _functions.tensor_times_scalar(h, 1 / v) + + if self._centering: + volume = dict_dimension["L"]**3 + mean = _functions.spatial_mean( + _functions.tensor_times_scalar(h, volume)) + h = h - mean + + h = self._mlp(phlower_tensor_collection({"h": h})) + assert h.dimension.is_dimensionless() + + if self._centering: + linear_mean = self._linear_weight( + phlower_tensor_collection({"mean": mean})) + h = h + linear_mean + + if self._invariant: + return h + + # Come back to the original dimension + for v in dict_scales.values(): + h = _functions.tensor_times_scalar(h, v) + + return h diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index f0dfd2d..d3761fc 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -16,6 +16,8 @@ ProportionalSetting, ) from phlower.settings._module_settings._share_setting import ShareSetting +from phlower.settings._module_settings \ + ._similarity_equivariant_mlp_setting import SimilarityEquivariantMLPSetting _name_to_setting: dict[str, IPhlowerLayerParameters] = { "Concatenator": ConcatenatorSetting, @@ -26,6 +28,7 @@ "PInvMLP": PInvMLPSetting, "Proportional": ProportionalSetting, "Share": ShareSetting, + "SimilarityEquivariantMLP": SimilarityEquivariantMLPSetting, } diff --git a/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py new file mode 100644 index 0000000..7c955e1 --- /dev/null +++ b/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import pydantic +from pydantic import Field +from typing_extensions import Self + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) + + +class SimilarityEquivariantMLPSetting( + IPhlowerLayerParameters, pydantic.BaseModel): + nodes: list[int] = Field( + ... + ) # This property only overwritten when resolving. + activations: list[str] = Field(default_factory=lambda: [], frozen=True) + dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) + bias: bool = Field(False, frozen=True) + create_linear_weight: bool = Field(False, frozen=True) + norm_function_name: str = Field( + default_factory=lambda: 'identity', frozen=True) + disable_en_equivariance: bool = Field(False, frozen=True) + invariant: bool = Field(False, frozen=True) + centering: bool = Field(False, frozen=True) + + def gather_input_dims(self, *input_dims: int) -> int: + if len(input_dims) != 1: + raise ValueError( + "Only one input is allowed in SimilarityEquivariantMLP.") + return input_dims[0] + + @pydantic.field_validator("nodes") + @classmethod + def check_n_nodes(cls, vals: list[int]) -> list[int]: + if len(vals) < 2: + raise ValueError( + "size of nodes must be larger than 1 in " + "SimilarityEquivariantMLPSetting. input: {vals}" + ) + + for i, v in enumerate(vals): + if v > 0: + continue + + if (i == 0) and (v == -1): + continue + + raise ValueError( + "nodes in SimilarityEquivariantMLPSetting is inconsistent. " + f"value {v} in {i}-th of nodes is not allowed." + ) + + return vals + + @pydantic.model_validator(mode="after") + def check_nodes_size(self) -> Self: + if len(self.nodes) - 1 != len(self.activations): + raise ValueError( + "Size of nodes and activations is not compatible " + "in SimilarityEquivariantMLPSetting." + " len(nodes) must be equal to 1 + len(activations)." + ) + return self + + def get_n_nodes(self) -> list[int] | None: + return self.nodes + + def overwrite_nodes(self, nodes: list[int]) -> None: + self.nodes = nodes + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + return diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 70a78b7..238a245 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -77,3 +77,11 @@ class PhlowerTypeError(TypeError): class PhlowerInvalidActivationError(ValueError): """This error raises when a set activation is invalid""" + +class PhlowerDimensionRequiredError(ValueError): + """ + This error raises when the dimension does not exist despite required. + """ + +class PhlowerInvalidArgumentsError(ValueError): + """This error raises when the arguments are invalid.""" diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index 4292e55..3b1230b 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -35,6 +35,19 @@ def test__add_with_unit(): np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c) +def test__sub_with_unit(): + units = phlower_dimension_tensor({"L": 2, "T": -2}) + a = np.random.rand(3, 10) + b = np.random.rand(3, 10) + c = a - b + + ap = PhlowerTensor(torch.tensor(a), units) + bp = PhlowerTensor(torch.tensor(b), units) + cp = ap - bp + + np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c) + + @pytest.mark.parametrize( "unit1, unit2", [({"L": 2, "T": -2}, None), (None, {"M": 2, "T": -3})] ) @@ -76,6 +89,53 @@ def test__mul_with_unit(): assert cp._dimension_tensor == dims_3 +def test__div_with_unit(): + dims_1 = phlower_dimension_tensor({"L": 2, "T": -2}) + dims_2 = phlower_dimension_tensor({"M": 1, "T": -2}) + dims_3 = phlower_dimension_tensor({"L": 2, "M": -1, "T": 0}) + + a = torch.tensor(np.random.rand(3, 10)) + b = torch.tensor(np.random.rand(3, 10)) + c = a / b + + ap = PhlowerTensor(a, dims_1) + bp = PhlowerTensor(b, dims_2) + cp = ap / bp + + np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) + + assert cp._dimension_tensor == dims_3 + + +def test__tensor_div_scalar(): + dims = phlower_dimension_tensor({"L": 2, "T": -2}) + + a = torch.tensor(np.random.rand(3, 10)) + c = a / 3. + + ap = PhlowerTensor(a, dims) + cp = ap / 3. + + np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) + + assert cp._dimension_tensor == dims + + +def test__scalar_div_tensor(): + dims = phlower_dimension_tensor({"L": 2, "T": -2}) + desired_dims = phlower_dimension_tensor({"L": -2, "T": 2}) + + a = torch.tensor(np.random.rand(3, 10)) + c = 3. / a + + ap = PhlowerTensor(a, dims) + cp = 3. / ap + + np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) + + assert cp._dimension_tensor == desired_dims + + def test__tanh(): a = torch.tensor(np.random.rand(3, 10)) c = np.tanh(a) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 892ac0c..023698d 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -572,3 +572,57 @@ def test_apply_orthogonal_group( np.testing.assert_almost_equal(actual_dimension, desired_dimension) else: assert actual.dimension is None + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel, mean_dims", + [ + ((10, 1), False, False, [0]), + ((10, 16), False, False, [0]), + ((10, 3, 16), False, False, [0]), + ((10, 3, 3, 16), False, False, [0]), + ((4, 10, 1), True, False, [1]), + ((4, 10, 16), True, False, [1]), + ((4, 10, 3, 16), True, False, [1]), + ((4, 10, 3, 3, 16), True, False, [1]), + ((10, 10, 10, 1), False, True, [0, 1, 2]), + ((10, 10, 10, 16), False, True, [0, 1, 2]), + ((10, 10, 10, 3, 16), False, True, [0, 1, 2]), + ((10, 10, 10, 3, 3, 16), False, True, [0, 1, 2]), + ((4, 10, 10, 10, 1), True, True, [1, 2, 3]), + ((4, 10, 10, 10, 16), True, True, [1, 2, 3]), + ((4, 10, 10, 10, 3, 16), True, True, [1, 2, 3]), + ((4, 10, 10, 10, 3, 3, 16), True, True, [1, 2, 3]), + ], +) +@pytest.mark.parametrize( + "dimension", + [ + None, + [[-1], [2], [0], [0], [0], [0], [0]], + [[1], [0], [1], [0], [0], [0], [0]], + [[-1], [-1], [2], [0], [1], [0], [0]], + ], +) +def test_spatial_mean( + size, is_time_series, is_voxel, mean_dims, dimension): + torch_tensor = torch.rand(*size) + x = phlower_tensor( + torch_tensor, dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + actual = _functions.spatial_mean(x) + + desired = torch_tensor + for dim in mean_dims: + desired = torch.mean(desired, dim=dim, keepdim=True) + np.testing.assert_almost_equal( + actual.to_numpy(), desired.numpy(), decimal=5) + + assert actual.is_time_series == is_time_series + assert actual.is_voxel == is_voxel + + if dimension is not None: + actual_dimension = actual.dimension._tensor.numpy() + np.testing.assert_almost_equal(actual_dimension, dimension) + else: + assert actual.dimension is None diff --git a/tests/test_nn/test_core_modules/test_pinv.py b/tests/test_nn/test_core_modules/test_pinv.py index 1697f03..bd8b25b 100644 --- a/tests/test_nn/test_core_modules/test_pinv.py +++ b/tests/test_nn/test_core_modules/test_pinv.py @@ -20,9 +20,9 @@ def test__can_call_parameters(): @pytest.mark.parametrize( "mlp_nodes, activations, decimal", [ - ([10, 10], ["identity"], 5), - ([10, 12], ["leaky_relu0p5"], 5), - ([20, 40, 100], ["tanh", "identity"], 5), + ([10, 10], ["identity"], 4), + ([10, 12], ["leaky_relu0p5"], 4), + ([20, 40, 100], ["tanh", "identity"], 4), ([20, 20, 40, 100], ["tanh", "smooth_leaky_relu", "leaky_relu0p5"], 3), ], ) diff --git a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py new file mode 100644 index 0000000..836bc43 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py @@ -0,0 +1,155 @@ +import numpy as np +import pytest +import torch +from scipy.stats import ortho_group + +from phlower import PhlowerTensor, phlower_tensor +from phlower._base._dimension import PhysicalDimensions +from phlower.collections import phlower_tensor_collection +from phlower.nn import SimilarityEquivariantMLP +from phlower.nn._core_modules import _functions +from phlower.utils.exceptions import ( + PhlowerDimensionRequiredError, + PhlowerInvalidArgumentsError, +) + + +def test__can_call_parameters(): + model = SimilarityEquivariantMLP(nodes=[8, 8]) + + # To check Concatenator inherit torch.nn.Module appropriately + _ = model.parameters() + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel", + [ + ((10, 1), False, False), + ((10, 16), False, False), + ((10, 3, 16), False, False), + ((4, 10, 1), True, False), + ((4, 10, 16), True, False), + ((4, 10, 3, 16), True, False), + ((10, 10, 10, 1), False, True), + ((10, 10, 10, 16), False, True), + ((10, 10, 10, 3, 16), False, True), + ((4, 10, 10, 10, 1), True, True), + ((4, 10, 10, 10, 16), True, True), + ((4, 10, 10, 10, 3, 16), True, True), + ], +) +@pytest.mark.parametrize("activation", ["identity", "tanh", "leaky_relu0p5"]) +@pytest.mark.parametrize("n_output_feature", [1, 16, 32]) +@pytest.mark.parametrize( + "dimension", + [ + # Dimensionless + {"T": 0, "L": 0, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Velocity + {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Mass density + {"T": 0, "L": - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Momentum density + {"T": -1, "L": 1 - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + ], +) +@pytest.mark.parametrize( + "norm_function_name", ["identity", "sqrt"]) +@pytest.mark.parametrize( + "disable_en_equivariance", [False, True]) +@pytest.mark.parametrize( + "centering", [False, True]) +@pytest.mark.parametrize( + "invariant", [False, True]) +def test__similarity_equivariance( + size, is_time_series, is_voxel, activation, n_output_feature, + dimension, norm_function_name, disable_en_equivariance, + centering, invariant, +): + if disable_en_equivariance: + orthogonal_tensor = PhlowerTensor(torch.eye(3)) + else: + orthogonal_tensor = PhlowerTensor( + torch.tensor(ortho_group.rvs(3).astype(np.float32))) + dict_scaling_factor = { + k: + np.random.rand() * 10 for k, v in dimension.items()} + scaling_factor = np.prod( + [dict_scaling_factor[k]**v for k, v in dimension.items()]) + + create_linear_weight = size[-1] != n_output_feature + model = SimilarityEquivariantMLP( + nodes=[size[-1], n_output_feature, n_output_feature], + activations=[activation, activation], + create_linear_weight=create_linear_weight, + norm_function_name=norm_function_name, + disable_en_equivariance=disable_en_equivariance, + centering=centering, invariant=invariant, + ) + + t = phlower_tensor( + torch.randn(*size), dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + dict_tensor = {"tensor": t} + scales = { + k: phlower_tensor( + torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) + for k, v in dimension.items() + if v != 0 + } + if len(scales) == 0: + scales = {"L": phlower_tensor( + torch.rand(1) * 10, dimension=PhysicalDimensions({"L": 1}))} + dict_tensor.update(scales) + + ts = phlower_tensor_collection(dict_tensor) + if invariant: + actual_tensor = _functions.apply_orthogonal_group( + orthogonal_tensor, model(ts)) * 1. + else: + actual_tensor = _functions.apply_orthogonal_group( + orthogonal_tensor, model(ts)) * scaling_factor + actual = actual_tensor.to_numpy() + + dict_transformed = {'tensor': _functions.apply_orthogonal_group( + orthogonal_tensor, t) * scaling_factor} + dict_scaled_scales = { + k: v * dict_scaling_factor[k] for k, v in scales.items()} + dict_transformed.update(dict_scaled_scales) + + transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) + desired = model(transformed_phlower_tensors) + + norm = torch.mean(_functions.contraction(desired)**.5).to_numpy() + # Test equivariance + np.testing.assert_almost_equal( + actual / norm, desired.to_numpy() / norm, decimal=2) + + if invariant: + # Test dimensionless in case of invariant + for v in actual_tensor.dimension.to_dict().values(): + np.testing.assert_almost_equal(v, 0.) + else: + # Test dimension is kept + for k, v in actual_tensor.dimension.to_dict().items(): + np.testing.assert_almost_equal(v, dimension[k]) + + +def test__similarity_equivariance_no_dimension(): + model = SimilarityEquivariantMLP(nodes=[8, 8]) + + t = phlower_tensor(torch.rand(10, 3, 8), dimension=None) + t_scale = phlower_tensor( + torch.tensor(0.1), dimension=PhysicalDimensions({"T": 1})) + ts = phlower_tensor_collection({"tensor": t, "T": t_scale}) + with pytest.raises(PhlowerDimensionRequiredError): + model(ts) + +def test__similarity_equivariance_no_scale_input(): + model = SimilarityEquivariantMLP(nodes=[8, 8]) + dimension = {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0} + + t = phlower_tensor(torch.rand(10, 3, 8), dimension=dimension) + ts = phlower_tensor_collection({"tensor": t}) + with pytest.raises(PhlowerInvalidArgumentsError): + model(ts) From d940fce63b4958f86121b8bc7ec61bbe09b634b7 Mon Sep 17 00:00:00 2001 From: horiem Date: Mon, 12 Aug 2024 03:15:21 +0900 Subject: [PATCH 31/89] update test --- .../test_core_modules/test_similarity_equivariant_mlp.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py index 836bc43..3508e7c 100644 --- a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py +++ b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py @@ -118,12 +118,11 @@ def test__similarity_equivariance( dict_transformed.update(dict_scaled_scales) transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) - desired = model(transformed_phlower_tensors) + desired = model(transformed_phlower_tensors).to_numpy() - norm = torch.mean(_functions.contraction(desired)**.5).to_numpy() + scale = np.max(desired) # Test equivariance - np.testing.assert_almost_equal( - actual / norm, desired.to_numpy() / norm, decimal=2) + np.testing.assert_almost_equal(actual / scale, desired / scale, decimal=2) if invariant: # Test dimensionless in case of invariant From 43c158fda262f8c207ac4f1da1443957ada3e172 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Fri, 9 Aug 2024 23:49:42 +0900 Subject: [PATCH 32/89] add all optimizers and schedulers defined in PyTroch --- .../loss_operations/_loss_calculator.py | 2 +- src/phlower/services/trainer/_optimizer.py | 51 ++++++ src/phlower/services/trainer/_trainer.py | 12 +- src/phlower/settings/__init__.py | 2 +- src/phlower/settings/_phlower_setting.py | 74 +-------- src/phlower/settings/_trainer_setting.py | 146 ++++++++++++++++++ src/phlower/utils/__init__.py | 2 + src/phlower/utils/_optimizer.py | 36 +++++ src/phlower/utils/_schedulers.py | 39 +++++ tests/e2e_tests/data/train.yml | 6 +- tests/e2e_tests/data/train_batch_size.yml | 5 +- tests/test_settings/test_phlower_setting.py | 3 + tests/test_utils/test_optimizer.py | 59 +++++++ tests/test_utils/test_schedulers.py | 67 ++++++++ 14 files changed, 422 insertions(+), 82 deletions(-) create mode 100644 src/phlower/services/trainer/_optimizer.py create mode 100644 src/phlower/settings/_trainer_setting.py create mode 100644 src/phlower/utils/_optimizer.py create mode 100644 src/phlower/utils/_schedulers.py create mode 100644 tests/test_utils/test_optimizer.py create mode 100644 tests/test_utils/test_schedulers.py diff --git a/src/phlower/services/loss_operations/_loss_calculator.py b/src/phlower/services/loss_operations/_loss_calculator.py index 1170008..0a4bb0d 100644 --- a/src/phlower/services/loss_operations/_loss_calculator.py +++ b/src/phlower/services/loss_operations/_loss_calculator.py @@ -12,7 +12,7 @@ phlower_tensor_collection, ) from phlower.services.loss_operations._loss_functions import get_loss_function -from phlower.settings._phlower_setting import LossSetting, PhlowerTrainerSetting +from phlower.settings._trainer_setting import LossSetting, PhlowerTrainerSetting from phlower.utils.typing import LossFunctionType diff --git a/src/phlower/services/trainer/_optimizer.py b/src/phlower/services/trainer/_optimizer.py new file mode 100644 index 0000000..9f0725c --- /dev/null +++ b/src/phlower/services/trainer/_optimizer.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import torch + +from phlower.settings import PhlowerTrainerSetting +from phlower.utils import OptimizerSelector, SchedulerSelector + + +class PhlowerOptimizerWrapper: + @classmethod + def from_setting( + cls, setting: PhlowerTrainerSetting, model: torch.nn.Module + ) -> PhlowerOptimizerWrapper: + return PhlowerOptimizerWrapper( + parameters=model.parameters(), + optimizer=setting.optimizer_setting.optimizer, + optimizer_kwards=setting.optimizer_setting.parameters, + schedulers={ + v.scheduler: v.parameters for v in setting.scheduler_setting + }, + ) + + def __init__( + self, + parameters: Iterator[torch.nn.Parameter], + optimizer: str, + optimizer_kwards: dict, + schedulers: dict[str, dict], + ): + self._optimizer = OptimizerSelector.select(optimizer)( + parameters, **optimizer_kwards + ) + self._schedulers = [ + SchedulerSelector.select(k)(self._optimizer, **v) + for k, v in schedulers.items() + ] + + def zero_grad(self): + self._optimizer.zero_grad() + + def step_optimizer(self): + self._optimizer.step() + + def step_scheduler(self): + for scheduler in self._schedulers: + scheduler.step() + + def state_dict(self): + return self._optimizer.state_dict() diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index 4abfec6..6e61f4e 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -11,6 +11,7 @@ from phlower.io import PhlowerCheckpointFile, PhlowerYamlFile from phlower.nn import PhlowerGroupModule from phlower.services.loss_operations import LossCalculator +from phlower.services.trainer._optimizer import PhlowerOptimizerWrapper from phlower.services.trainer._trainer_logger import LogRecord, LogRecordIO from phlower.settings import ( PhlowerModelSetting, @@ -58,8 +59,8 @@ def __init__( self._model_setting.network ) self._model.to(self._trainer_setting.device) - self._optimizer = torch.optim.SGD( - self._model.parameters(), lr=self._trainer_setting.lr, momentum=0.9 + self._scheduled_optimizer = PhlowerOptimizerWrapper.from_setting( + self._trainer_setting, model=self._model ) self._timer = StopWatch() @@ -121,7 +122,7 @@ def train( self._model.train() for tr_batch in train_loader: tr_batch: LumpedTensorData - self._optimizer.zero_grad() + self._scheduled_optimizer.zero_grad() h = self._model.forward( tr_batch.x_data, supports=tr_batch.sparse_supports @@ -133,12 +134,13 @@ def train( loss = loss_function.aggregate(losses) train_losses.append(loss.detach().to_tensor().float().item()) loss.backward() - self._optimizer.step() + self._scheduled_optimizer.step_optimizer() _train_batch_pbar.update( trick=self._trainer_setting.batch_size, desc=f"batch train loss: {train_losses[-1]:.3f}", ) + self._scheduled_optimizer.step_scheduler() self._model.eval() for val_batch in validation_loader: @@ -216,7 +218,7 @@ def _save_checkpoint( "epoch": epoch, "validation_loss": validation_loss, "model_state_dict": self._model.state_dict(), - "optimizer_state_dict": self._optimizer.state_dict(), + "optimizer_state_dict": self._scheduled_optimizer.state_dict(), } prefix = PhlowerCheckpointFile.get_fixed_prefix() file_basename = f"{prefix}{epoch}" diff --git a/src/phlower/settings/__init__.py b/src/phlower/settings/__init__.py index 8100946..cd3f3f1 100644 --- a/src/phlower/settings/__init__.py +++ b/src/phlower/settings/__init__.py @@ -3,10 +3,10 @@ PhlowerModelSetting, PhlowerPredictorSetting, PhlowerSetting, - PhlowerTrainerSetting, ) from phlower.settings._scaling_setting import ( PhlowerScalingSetting, ScalerInputParameters, ScalerResolvedParameter, ) +from phlower.settings._trainer_setting import PhlowerTrainerSetting diff --git a/src/phlower/settings/_phlower_setting.py b/src/phlower/settings/_phlower_setting.py index db4a6ad..d926ec2 100644 --- a/src/phlower/settings/_phlower_setting.py +++ b/src/phlower/settings/_phlower_setting.py @@ -1,7 +1,6 @@ from __future__ import annotations import pathlib -from collections.abc import Iterable import pydantic from pydantic import dataclasses as dc @@ -11,6 +10,7 @@ from phlower.io import PhlowerYamlFile from phlower.settings._group_settings import GroupModuleSetting from phlower.settings._scaling_setting import PhlowerScalingSetting +from phlower.settings._trainer_setting import PhlowerTrainerSetting from phlower.utils.enums import ModelSelectionType @@ -81,78 +81,6 @@ class PhlowerModelSetting(pydantic.BaseModel): ) -class PhlowerTrainerSetting(pydantic.BaseModel): - loss_setting: LossSetting - """ - setting for loss function - """ - - n_epoch: int = 10 - """ - the number of epochs. Defaults to 10. - """ - - random_seed: int = 0 - """ - random seed. Defaults to 0 - """ - - lr: float = 0.001 - """ - learning rate. Defaults to 0.001 - """ - - batch_size: int = 1 - """ - batch size. Defaults to 1 - """ - - num_workers: int = 1 - """ - the number of cores. Defaults to 1. - """ - - device: str = "cpu" - """ - device name. Defaults to cpu - """ - - non_blocking: bool = False - - # special keyward to forbid extra fields in pydantic - model_config = pydantic.ConfigDict(frozen=True, extra="forbid") - - -@dc.dataclass(frozen=True, config=pydantic.ConfigDict(extra="forbid")) -class LossSetting: - name2loss: dict[str, str] - """ - Dictionary which maps name of target variable to name of loss function. - """ - - name2weight: dict[str, str] | None = None - """ - Dictionary which maps weight value to name of output variable. - Defaults to None. If None, total loss value is calculated as summation of all loss values. - """ - - def loss_names(self) -> Iterable[str]: - """get registered loss names - - Returns: - Iterable[str]: names of loss functions - """ - return self.name2loss.values() - - def loss_variable_names(self) -> list[str]: - """get variable names - - Returns: - list[str]: variable names - """ - return list(self.name2loss.keys()) - - @dc.dataclass(frozen=True, config=pydantic.ConfigDict(extra="forbid")) class PhlowerPredictorSetting: selection_mode: str diff --git a/src/phlower/settings/_trainer_setting.py b/src/phlower/settings/_trainer_setting.py new file mode 100644 index 0000000..398b591 --- /dev/null +++ b/src/phlower/settings/_trainer_setting.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +from collections.abc import Iterable + +import pydantic +import pydantic.dataclasses as dc +from pydantic import Field + +from phlower.utils import OptimizerSelector, SchedulerSelector + + +@dc.dataclass(frozen=True, config=pydantic.ConfigDict(extra="forbid")) +class LossSetting: + name2loss: dict[str, str] + """ + Dictionary which maps name of target variable to name of loss function. + """ + + name2weight: dict[str, str] | None = None + """ + Dictionary which maps weight value to name of output variable. + Defaults to None. If None, total loss value is calculated as summation of all loss values. + """ + + def loss_names(self) -> Iterable[str]: + """get registered loss names + + Returns: + Iterable[str]: names of loss functions + """ + return self.name2loss.values() + + def loss_variable_names(self) -> list[str]: + """get variable names + + Returns: + list[str]: variable names + """ + return list(self.name2loss.keys()) + + +class OptimizerSetting(pydantic.BaseModel): + optimizer: str = "Adam" + """ + Optimizer Class name defined in torch.optim. Default to Adam. + Ex. Adam, RMSprop, SGD + """ + + parameters: dict[str, int | float | bool | str] = Field( + default_factory=dict + ) + """ + Parameters to pass when optimizer class is initialized. + Allowed parameters depend on the optimizer you choose. + """ + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(frozen=True, extra="forbid") + + @pydantic.field_validator("optimizer") + @classmethod + def check_exist_scheduler(cls, name): + if not OptimizerSelector.exist(name): + raise ValueError( + f"{name} is not defined in phlower. " + "If you defined user defined optimizer, " + "please use `register` function in " + "`phlower.utils.OptimizerSelector`." + ) + return name + + +class SchedulerSetting(pydantic.BaseModel): + scheduler: str + """ + Scheduler Class name defined in torch.optim.lr_schedulers. + """ + + parameters: dict[str, int | float | bool | str] = Field( + default_factory=dict + ) + """ + Parameters to pass when scheduler class is initialized. + Allowed parameters depend on the scheduler you choose. + """ + + @pydantic.field_validator("scheduler") + @classmethod + def check_exist_scheduler(cls, name): + if not SchedulerSelector.exist(name): + raise ValueError( + f"{name} is not defined in phlower. " + "If you defined user defined scheduler, " + "please use `register` function in " + "`phlower.utils.SchedulerSelector`." + ) + return name + + +class PhlowerTrainerSetting(pydantic.BaseModel): + loss_setting: LossSetting + """ + setting for loss function + """ + + optimizer_setting: OptimizerSetting = Field( + default_factory=OptimizerSetting + ) + """ + setting for optimizer + """ + + scheduler_setting: list[SchedulerSetting] = Field(default_factory=list) + """ + setting for schedulers + """ + + n_epoch: int = 10 + """ + the number of epochs. Defaults to 10. + """ + + random_seed: int = 0 + """ + random seed. Defaults to 0 + """ + + batch_size: int = 1 + """ + batch size. Defaults to 1 + """ + + num_workers: int = 1 + """ + the number of cores. Defaults to 1. + """ + + device: str = "cpu" + """ + device name. Defaults to cpu + """ + + non_blocking: bool = False + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(frozen=True, extra="forbid") diff --git a/src/phlower/utils/__init__.py b/src/phlower/utils/__init__.py index c658eb9..f7139f0 100644 --- a/src/phlower/utils/__init__.py +++ b/src/phlower/utils/__init__.py @@ -8,3 +8,5 @@ from phlower.utils._progress_bar import PhlowerProgressBar from phlower.utils._multiprocessor import PhlowerMultiprocessor from phlower.utils._timer import StopWatch +from phlower.utils._optimizer import OptimizerSelector +from phlower.utils._schedulers import SchedulerSelector diff --git a/src/phlower/utils/_optimizer.py b/src/phlower/utils/_optimizer.py new file mode 100644 index 0000000..85b5c73 --- /dev/null +++ b/src/phlower/utils/_optimizer.py @@ -0,0 +1,36 @@ +import torch + + +class OptimizerSelector: + _REGISTERED: dict[str, type[torch.optim.Optimizer]] = { + "Adadelta": torch.optim.Adadelta, + "Adam": torch.optim.Adam, + "AdamW": torch.optim.AdamW, + "SparseAdam": torch.optim.SparseAdam, + "Adamax": torch.optim.Adamax, + "ASGD": torch.optim.ASGD, + "LBFGS": torch.optim.LBFGS, + "NAdam": torch.optim.NAdam, + "RAdam": torch.optim.RAdam, + "RMSprop": torch.optim.RMSprop, + "Rprop": torch.optim.Rprop, + "SGD": torch.optim.SGD, + } + + @staticmethod + def register(name: str, cls: type[torch.optim.Optimizer]) -> None: + OptimizerSelector._REGISTERED.update({name: cls}) + + @staticmethod + def exist(name: str): + return name in OptimizerSelector._REGISTERED + + @staticmethod + def select(name: str) -> type[torch.optim.Optimizer]: + _dict = OptimizerSelector._REGISTERED + if name not in _dict: + raise KeyError( + f"{name} is not implemented as an optimizer in phlower. " + f"Registered optimizers are {list(_dict.keys())}" + ) + return _dict[name] diff --git a/src/phlower/utils/_schedulers.py b/src/phlower/utils/_schedulers.py new file mode 100644 index 0000000..dc73350 --- /dev/null +++ b/src/phlower/utils/_schedulers.py @@ -0,0 +1,39 @@ +import torch + + +class SchedulerSelector: + _REGISTERED: dict[str, type[torch.optim.lr_scheduler.LRScheduler]] = { + "LambdaLR": torch.optim.lr_scheduler.LambdaLR, + "MultiplicativeLR": torch.optim.lr_scheduler.MultiplicativeLR, + "StepLR": torch.optim.lr_scheduler.StepLR, + "MultiStepLR": torch.optim.lr_scheduler.MultiStepLR, + "ConstantLR": torch.optim.lr_scheduler.ConstantLR, + "LinearLR": torch.optim.lr_scheduler.LinearLR, + "ExponentialLR": torch.optim.lr_scheduler.ExponentialLR, + "SequentialLR": torch.optim.lr_scheduler.SequentialLR, + "CosineAnnealingLR": torch.optim.lr_scheduler.CosineAnnealingLR, + "ChainedScheduler": torch.optim.lr_scheduler.ChainedScheduler, + "ReduceLROnPlateau": torch.optim.lr_scheduler.ReduceLROnPlateau, + "CyclicLR": torch.optim.lr_scheduler.CyclicLR, + "CosineAnnealingWarmRestarts": torch.optim.lr_scheduler.CosineAnnealingWarmRestarts, + "OneCycleLR": torch.optim.lr_scheduler.OneCycleLR, + "PolynomialLR": torch.optim.lr_scheduler.PolynomialLR, + } + + @staticmethod + def register(name: str, cls: type[torch.optim.Optimizer]) -> None: + SchedulerSelector._REGISTERED.update({name: cls}) + + @staticmethod + def exist(name: str): + return name in SchedulerSelector._REGISTERED + + @staticmethod + def select(name: str) -> type[torch.optim.Optimizer]: + _dict = SchedulerSelector._REGISTERED + if name not in _dict: + raise KeyError( + f"{name} is not implemented as an scheduler in phlower. " + f"Registered schedulers are {list(_dict.keys())}" + ) + return _dict[name] diff --git a/tests/e2e_tests/data/train.yml b/tests/e2e_tests/data/train.yml index ee0203b..9b41c62 100644 --- a/tests/e2e_tests/data/train.yml +++ b/tests/e2e_tests/data/train.yml @@ -2,10 +2,14 @@ training: batch_size: 1 random_seed: 0 - lr: 0.0001 loss_setting: name2loss: nodal_last_u: "mse" + optimizer_setting: + optimizer: SGD + parameters: + lr: 0.0001 + model: variable_dimensions: diff --git a/tests/e2e_tests/data/train_batch_size.yml b/tests/e2e_tests/data/train_batch_size.yml index 8c6b945..409c59e 100644 --- a/tests/e2e_tests/data/train_batch_size.yml +++ b/tests/e2e_tests/data/train_batch_size.yml @@ -2,7 +2,10 @@ training: batch_size: 3 random_seed: 0 - lr: 0.0001 + optimizer_setting: + optimizer: SGD + parameters: + lr: 0.0001 loss_setting: name2loss: nodal_last_u: "mse" diff --git a/tests/test_settings/test_phlower_setting.py b/tests/test_settings/test_phlower_setting.py index f262f15..ed32d00 100644 --- a/tests/test_settings/test_phlower_setting.py +++ b/tests/test_settings/test_phlower_setting.py @@ -1,4 +1,5 @@ import pathlib +import shutil import pydantic import pytest @@ -46,6 +47,8 @@ def test__model_dump(): "tests/test_settings/data/e2e/setting1.yml" ) + output_directory = pathlib.Path("tests/test_settings/tmp") + shutil.rmtree(output_directory, ignore_errors=True) _ = PhlowerYamlFile.save( output_directory=pathlib.Path("tests/test_settings/tmp"), file_basename="output", diff --git a/tests/test_utils/test_optimizer.py b/tests/test_utils/test_optimizer.py new file mode 100644 index 0000000..6e6e95b --- /dev/null +++ b/tests/test_utils/test_optimizer.py @@ -0,0 +1,59 @@ +from unittest import mock + +import pytest +import torch + +from phlower.utils import OptimizerSelector + + +@pytest.mark.parametrize( + "name, desired", + [ + ("Adadelta", True), + ("Adam", True), + ("SparseAdam", True), + ("Adamax", True), + ("ASGD", True), + ("LBFGS", True), + ("NAdam", True), + ("RAdam", True), + ("RMSprop", True), + ("Rprop", True), + ("SGD", True), + ("MyOptimizer", False), + ], +) +def test__exist(name, desired): + assert OptimizerSelector.exist(name) == desired + + +@pytest.mark.parametrize( + "name", + [ + ("Adadelta"), + ("Adam"), + ("SparseAdam"), + ("Adamax"), + ("ASGD"), + ("LBFGS"), + ("NAdam"), + ("RAdam"), + ("RMSprop"), + ("Rprop"), + ("SGD"), + ], +) +def test__select(name): + optim = OptimizerSelector.select(name) + assert optim.__name__ == name + + +@pytest.mark.parametrize("name", ["MyOptimizer", "BestOptimizer"]) +def test__exist_after_register(name): + assert not OptimizerSelector.exist(name) + dummy = mock.MagicMock(torch.optim.Optimizer) + OptimizerSelector.register(name, dummy) + assert OptimizerSelector.exist(name) + + # Not to affect other tests, unregister scheduler + OptimizerSelector._REGISTERED.pop(name) diff --git a/tests/test_utils/test_schedulers.py b/tests/test_utils/test_schedulers.py new file mode 100644 index 0000000..1adb1f4 --- /dev/null +++ b/tests/test_utils/test_schedulers.py @@ -0,0 +1,67 @@ +from unittest import mock + +import pytest +import torch + +from phlower.utils import SchedulerSelector + + +@pytest.mark.parametrize( + "name, desired", + [ + ("LambdaLR", True), + ("MultiplicativeLR", True), + ("StepLR", True), + ("MultiStepLR", True), + ("ConstantLR", True), + ("LinearLR", True), + ("ExponentialLR", True), + ("SequentialLR", True), + ("CosineAnnealingLR", True), + ("ChainedScheduler", True), + ("ReduceLROnPlateau", True), + ("CyclicLR", True), + ("CosineAnnealingWarmRestarts", True), + ("OneCycleLR", True), + ("PolynomialLR", True), + ("MyScheduler", False), + ], +) +def test__exist(name, desired): + assert SchedulerSelector.exist(name) == desired + + +@pytest.mark.parametrize( + "name", + [ + ("LambdaLR"), + ("MultiplicativeLR"), + ("StepLR"), + ("MultiStepLR"), + ("ConstantLR"), + ("LinearLR"), + ("ExponentialLR"), + ("SequentialLR"), + ("CosineAnnealingLR"), + ("ChainedScheduler"), + ("ReduceLROnPlateau"), + ("CyclicLR"), + ("CosineAnnealingWarmRestarts"), + ("OneCycleLR"), + ("PolynomialLR"), + ], +) +def test__select(name): + scheduler = SchedulerSelector.select(name) + assert scheduler.__name__ == name + + +@pytest.mark.parametrize("name", ["MyScheduler", "BestScheduler"]) +def test__exist_after_register(name): + assert not SchedulerSelector.exist(name) + dummy = mock.MagicMock(torch.optim.Optimizer) + SchedulerSelector.register(name, dummy) + assert SchedulerSelector.exist(name) + + # Not to affect other tests, unregister scheduler + SchedulerSelector._REGISTERED.pop(name) From 1cc52c31763a7f664f503359b5ce22da0d7a71da Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 13 Aug 2024 16:07:53 +0900 Subject: [PATCH 33/89] add tests for optimizers --- .../_base/array/dense/_ndarray_wrapper.py | 6 +- .../array/sparse/_sparse_array_wrapper.py | 7 +- src/phlower/_base/tensors/_phlower_tensor.py | 118 +++++++++++------ .../tensors/_unsupported_function_names.py | 5 +- .../collections/arrays/_arrays_dict.py | 7 +- src/phlower/nn/_core_modules/_functions.py | 12 +- src/phlower/services/trainer/_optimizer.py | 13 +- src/phlower/settings/_trainer_setting.py | 6 +- src/phlower/utils/_schedulers.py | 8 +- src/phlower/utils/exceptions.py | 1 + .../test_tensors/test__phlower_tensor.py | 97 +++++++------- .../test_core_modules/test_functions.py | 11 +- tests/test_nn/test_core_modules/test_gcn.py | 13 +- .../test_trainer/test_optimizer.py | 121 ++++++++++++++++++ 14 files changed, 308 insertions(+), 117 deletions(-) create mode 100644 tests/test_services/test_trainer/test_optimizer.py diff --git a/src/phlower/_base/array/dense/_ndarray_wrapper.py b/src/phlower/_base/array/dense/_ndarray_wrapper.py index c7fe74d..bcb160d 100644 --- a/src/phlower/_base/array/dense/_ndarray_wrapper.py +++ b/src/phlower/_base/array/dense/_ndarray_wrapper.py @@ -70,8 +70,10 @@ def to_phlower_tensor( if is_time_series is None: is_time_series = self.is_time_series _tensor = phlower_tensor( - tensor=torch.from_numpy(self.data), dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel, + tensor=torch.from_numpy(self.data), + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, ) _tensor.to(device=device, non_blocking=non_blocking) return _tensor diff --git a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py index ea8c75b..adef209 100644 --- a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py +++ b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py @@ -110,8 +110,11 @@ def to_phlower_tensor( self._sparse_data.shape, ) _tensor = phlower_tensor( - tensor=sparse_tensor, dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + tensor=sparse_tensor, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) _tensor = _tensor.coalesce() _tensor.to(device=device, non_blocking=non_blocking) return _tensor diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 73a2b6a..f7a0f12 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -50,8 +50,11 @@ def phlower_tensor( dimension_tensor = _resolve_dimension_arg(dimension) return PhlowerTensor( - tensor=tensor, dimension_tensor=dimension_tensor, - is_time_series=is_time_series, is_voxel=is_voxel) + tensor=tensor, + dimension_tensor=dimension_tensor, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) def _resolve_dimension_arg( @@ -157,7 +160,8 @@ def rank(self) -> int: """Returns the tensor rank.""" if self.is_sparse: raise PhlowerSparseUnsupportedError( - "Cannot call rank() for sparse PhlowerTensor") + "Cannot call rank() for sparse PhlowerTensor" + ) size = self.size() start = 1 if self.is_time_series: @@ -170,14 +174,15 @@ def n_vertices(self) -> int: """Returns the number of vertices.""" if self.is_sparse: raise PhlowerSparseUnsupportedError( - "Cannot call n_vertices() for sparse PhlowerTensor") + "Cannot call n_vertices() for sparse PhlowerTensor" + ) size = self.size() start = 0 if self.is_time_series: start += 1 if self.is_voxel: - return np.prod(size[start:start+3]) + return np.prod(size[start : start + 3]) return size[start] @@ -213,50 +218,72 @@ def to_vertexwise(self) -> PhlowerTensor: t_pattern = "" if self.is_voxel: space_pattern = "x y z " - dict_shape.update({ - "x": shape[space_start], - "y": shape[space_start + 1], - "z": shape[space_start + 2], - }) + dict_shape.update( + { + "x": shape[space_start], + "y": shape[space_start + 1], + "z": shape[space_start + 2], + } + ) feat_start = space_start + 3 else: space_pattern = "n " dict_shape.update({"n": shape[space_start]}) feat_start = space_start + 1 - feat_pattern = " ".join([ - f"a{i}" for i in range(len(shape[feat_start:]))]) + feat_pattern = " ".join( + [f"a{i}" for i in range(len(shape[feat_start:]))] + ) # Do not include the last axis in case modified by NNs. - dict_shape.update({ - f"a{i}": s for i, s in enumerate(shape[feat_start:-1])}) + dict_shape.update( + {f"a{i}": s for i, s in enumerate(shape[feat_start:-1])} + ) original_pattern = f"{t_pattern}{space_pattern}{feat_pattern}" resultant_pattern = f"({space_pattern}) ({t_pattern}{feat_pattern})" tensor_2d = einops.rearrange( - self.to_tensor(), f"{original_pattern} -> {resultant_pattern}") - return PhlowerTensor( - tensor_2d, dimension_tensor=self.dimension, - is_time_series=False, is_voxel=False), \ - original_pattern, resultant_pattern, dict_shape + self.to_tensor(), f"{original_pattern} -> {resultant_pattern}" + ) + return ( + PhlowerTensor( + tensor_2d, + dimension_tensor=self.dimension, + is_time_series=False, + is_voxel=False, + ), + original_pattern, + resultant_pattern, + dict_shape, + ) def rearrange( - self, pattern: str, - is_time_series: bool = False, is_voxel: bool = False, - **kwargs) -> PhlowerTensor: + self, + pattern: str, + is_time_series: bool = False, + is_voxel: bool = False, + **kwargs, + ) -> PhlowerTensor: tensor = self.to_tensor() rearranged = einops.rearrange(tensor, pattern, **kwargs) return PhlowerTensor( - rearranged, dimension_tensor=self.dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + rearranged, + dimension_tensor=self.dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) def reshape( - self, shape: Sequence[int], - is_time_series: bool = False, is_voxel: bool = False, + self, + shape: Sequence[int], + is_time_series: bool = False, + is_voxel: bool = False, ) -> PhlowerTensor: return PhlowerTensor( torch.reshape(self.to_tensor(), shape), dimension_tensor=self.dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + is_time_series=is_time_series, + is_voxel=is_voxel, + ) def to(self, device: str, non_blocking: bool = False) -> None: self._tensor.to(device, non_blocking=non_blocking) @@ -267,7 +294,8 @@ def detach(self) -> PhlowerTensor: return PhlowerTensor( self._tensor.detach(), dimension_tensor=self._dimension_tensor.detach(), - is_time_series=self.is_time_series, is_voxel=self.is_voxel, + is_time_series=self.is_time_series, + is_voxel=self.is_voxel, ) def backward(self) -> None: @@ -283,7 +311,8 @@ def __torch_function__( ): if func.__name__ in UNSUPPORTED_FUNCTION_NAMES: raise PhlowerUnsupportedTorchFunctionError( - f"Unsupported function: {func.__name__}") + f"Unsupported function: {func.__name__}" + ) if kwargs is None: kwargs = {} @@ -293,14 +322,15 @@ def __torch_function__( # NOTE: Assume flags for the first tensor is preserved is_time_series = _recursive_resolve( - args, "_is_time_series", return_first_only=True) - is_voxel = _recursive_resolve( - args, "_is_voxel", return_first_only=True) + args, "_is_time_series", return_first_only=True + ) + is_voxel = _recursive_resolve(args, "_is_voxel", return_first_only=True) if not _has_dimension(args): # Unit calculation is not considered when unit tensor is not found. return PhlowerTensor( - ret, is_time_series=is_time_series, is_voxel=is_voxel) + ret, is_time_series=is_time_series, is_voxel=is_voxel + ) _dimensions = _recursive_resolve( args, "_dimension_tensor", allow_none=False @@ -308,24 +338,32 @@ def __torch_function__( result_units = func(*_dimensions, **kwargs) return PhlowerTensor( - ret, result_units, is_time_series=is_time_series, - is_voxel=is_voxel) + ret, result_units, is_time_series=is_time_series, is_voxel=is_voxel + ) def _recursive_resolve( - args: Iterable | Any, attr: str, allow_none: bool = True, + args: Iterable | Any, + attr: str, + allow_none: bool = True, return_first_only: bool = False, ) -> list[str]: if isinstance(args, tuple | list): if return_first_only: return _recursive_resolve( - args[0], attr, allow_none=allow_none, - return_first_only=return_first_only) + args[0], + attr, + allow_none=allow_none, + return_first_only=return_first_only, + ) else: return [ _recursive_resolve( - v, attr, allow_none=allow_none, - return_first_only=return_first_only) + v, + attr, + allow_none=allow_none, + return_first_only=return_first_only, + ) for v in args ] diff --git a/src/phlower/_base/tensors/_unsupported_function_names.py b/src/phlower/_base/tensors/_unsupported_function_names.py index c2ef632..c17b45e 100644 --- a/src/phlower/_base/tensors/_unsupported_function_names.py +++ b/src/phlower/_base/tensors/_unsupported_function_names.py @@ -1,5 +1,4 @@ - UNSUPPORTED_FUNCTION_NAMES = [ - 'einsum', - 'reshape', + "einsum", + "reshape", ] diff --git a/src/phlower/collections/arrays/_arrays_dict.py b/src/phlower/collections/arrays/_arrays_dict.py index 6d735d5..8088114 100644 --- a/src/phlower/collections/arrays/_arrays_dict.py +++ b/src/phlower/collections/arrays/_arrays_dict.py @@ -63,8 +63,11 @@ def to_batched_tensor( ): tensors = [ v.to_phlower_tensor( - device=device, non_blocking=non_blocking, dimension=dimensions, - is_time_series=is_time_series, is_voxel=is_voxel, + device=device, + non_blocking=non_blocking, + dimension=dimensions, + is_time_series=is_time_series, + is_voxel=is_voxel, ) for v in self._data ] diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index cc572cd..6f67f3f 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -1,12 +1,11 @@ - import torch from phlower._base.tensors._interface import IPhlowerTensor def spmm( - sparse: IPhlowerTensor, x: IPhlowerTensor, - repeat: int = 1) -> IPhlowerTensor: + sparse: IPhlowerTensor, x: IPhlowerTensor, repeat: int = 1 +) -> IPhlowerTensor: """ Computes sparse matrix times dense tensor along with the vertex axis. @@ -26,5 +25,8 @@ def spmm( for _ in range(repeat): h = torch.sparse.mm(sparse, h) return h.rearrange( - pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, - **dict_shape) + pattern, + is_time_series=x.is_time_series, + is_voxel=x.is_voxel, + **dict_shape, + ) diff --git a/src/phlower/services/trainer/_optimizer.py b/src/phlower/services/trainer/_optimizer.py index 9f0725c..61e8192 100644 --- a/src/phlower/services/trainer/_optimizer.py +++ b/src/phlower/services/trainer/_optimizer.py @@ -16,7 +16,7 @@ def from_setting( return PhlowerOptimizerWrapper( parameters=model.parameters(), optimizer=setting.optimizer_setting.optimizer, - optimizer_kwards=setting.optimizer_setting.parameters, + optimizer_kwargs=setting.optimizer_setting.parameters, schedulers={ v.scheduler: v.parameters for v in setting.scheduler_setting }, @@ -26,11 +26,16 @@ def __init__( self, parameters: Iterator[torch.nn.Parameter], optimizer: str, - optimizer_kwards: dict, - schedulers: dict[str, dict], + optimizer_kwargs: dict | None = None, + schedulers: dict[str, dict] | None = None, ): + if optimizer_kwargs is None: + optimizer_kwargs = {} + if schedulers is None: + schedulers = {} + self._optimizer = OptimizerSelector.select(optimizer)( - parameters, **optimizer_kwards + parameters, **optimizer_kwargs ) self._schedulers = [ SchedulerSelector.select(k)(self._optimizer, **v) diff --git a/src/phlower/settings/_trainer_setting.py b/src/phlower/settings/_trainer_setting.py index 398b591..4b72a28 100644 --- a/src/phlower/settings/_trainer_setting.py +++ b/src/phlower/settings/_trainer_setting.py @@ -19,7 +19,8 @@ class LossSetting: name2weight: dict[str, str] | None = None """ Dictionary which maps weight value to name of output variable. - Defaults to None. If None, total loss value is calculated as summation of all loss values. + Defaults to None. If None, total loss value is calculated + as summation of all loss values. """ def loss_names(self) -> Iterable[str]: @@ -84,6 +85,9 @@ class SchedulerSetting(pydantic.BaseModel): Allowed parameters depend on the scheduler you choose. """ + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(frozen=True, extra="forbid") + @pydantic.field_validator("scheduler") @classmethod def check_exist_scheduler(cls, name): diff --git a/src/phlower/utils/_schedulers.py b/src/phlower/utils/_schedulers.py index dc73350..097a8ae 100644 --- a/src/phlower/utils/_schedulers.py +++ b/src/phlower/utils/_schedulers.py @@ -15,13 +15,15 @@ class SchedulerSelector: "ChainedScheduler": torch.optim.lr_scheduler.ChainedScheduler, "ReduceLROnPlateau": torch.optim.lr_scheduler.ReduceLROnPlateau, "CyclicLR": torch.optim.lr_scheduler.CyclicLR, - "CosineAnnealingWarmRestarts": torch.optim.lr_scheduler.CosineAnnealingWarmRestarts, + "CosineAnnealingWarmRestarts": torch.optim.lr_scheduler.CosineAnnealingWarmRestarts, # NOQA "OneCycleLR": torch.optim.lr_scheduler.OneCycleLR, "PolynomialLR": torch.optim.lr_scheduler.PolynomialLR, } @staticmethod - def register(name: str, cls: type[torch.optim.Optimizer]) -> None: + def register( + name: str, cls: type[torch.optim.lr_scheduler.LRScheduler] + ) -> None: SchedulerSelector._REGISTERED.update({name: cls}) @staticmethod @@ -29,7 +31,7 @@ def exist(name: str): return name in SchedulerSelector._REGISTERED @staticmethod - def select(name: str) -> type[torch.optim.Optimizer]: + def select(name: str) -> type[torch.optim.lr_scheduler.LRScheduler]: _dict = SchedulerSelector._REGISTERED if name not in _dict: raise KeyError( diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 9dbe5f6..d77832d 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -63,4 +63,5 @@ class PhlowerReshapeError(ValueError): manner """ + class NotFoundReferenceModuleError(ValueError): ... diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index 4292e55..c4d5358 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -92,21 +92,22 @@ def test__tanh(): (False, False, [100, 16], 0), (False, False, [100, 3, 16], 1), (False, False, [100, 3, 3, 16], 2), - ( True, False, [4, 100, 16], 0), - ( True, False, [4, 100, 3, 16], 1), - ( True, False, [4, 100, 3, 3, 16], 2), - (False, True, [10, 10, 10, 16], 0), - (False, True, [10, 10, 10, 3, 16], 1), - (False, True, [10, 10, 10, 3, 3, 16], 2), - ( True, True, [4, 10, 10, 10, 16], 0), - ( True, True, [4, 10, 10, 10, 3, 16], 1), - ( True, True, [4, 10, 10, 10, 3, 3, 16], 2), + (True, False, [4, 100, 16], 0), + (True, False, [4, 100, 3, 16], 1), + (True, False, [4, 100, 3, 3, 16], 2), + (False, True, [10, 10, 10, 16], 0), + (False, True, [10, 10, 10, 3, 16], 1), + (False, True, [10, 10, 10, 3, 3, 16], 2), + (True, True, [4, 10, 10, 10, 16], 0), + (True, True, [4, 10, 10, 10, 3, 16], 1), + (True, True, [4, 10, 10, 10, 3, 3, 16], 2), ], ) def test__rank(is_time_series, is_voxel, size, desired_rank): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( - torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel + ) assert phlower_tensor.rank() == desired_rank @@ -116,21 +117,22 @@ def test__rank(is_time_series, is_voxel, size, desired_rank): (False, False, [100, 16], 100), (False, False, [100, 3, 16], 100), (False, False, [100, 3, 3, 16], 100), - ( True, False, [4, 100, 16], 100), - ( True, False, [4, 100, 3, 16], 100), - ( True, False, [4, 100, 3, 3, 16], 100), - (False, True, [10, 10, 10, 16], 1000), - (False, True, [10, 10, 10, 3, 16], 1000), - (False, True, [10, 10, 10, 3, 3, 16], 1000), - ( True, True, [4, 10, 10, 10, 16], 1000), - ( True, True, [4, 10, 10, 10, 3, 16], 1000), - ( True, True, [4, 10, 10, 10, 3, 3, 16], 1000), + (True, False, [4, 100, 16], 100), + (True, False, [4, 100, 3, 16], 100), + (True, False, [4, 100, 3, 3, 16], 100), + (False, True, [10, 10, 10, 16], 1000), + (False, True, [10, 10, 10, 3, 16], 1000), + (False, True, [10, 10, 10, 3, 3, 16], 1000), + (True, True, [4, 10, 10, 10, 16], 1000), + (True, True, [4, 10, 10, 10, 3, 16], 1000), + (True, True, [4, 10, 10, 10, 3, 3, 16], 1000), ], ) def test__n_vertices(is_time_series, is_voxel, size, desired_n_vertices): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( - torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel + ) assert phlower_tensor.n_vertices() == desired_n_vertices @@ -147,21 +149,22 @@ def test__raises_phlower_sparse_rank_undefined_error(): (False, False, [100, 16], (100, 16)), (False, False, [100, 3, 16], (100, 3 * 16)), (False, False, [100, 3, 3, 16], (100, 3 * 3 * 16)), - ( True, False, [4, 100, 16], (100, 4 * 16)), - ( True, False, [4, 100, 3, 16], (100, 4 * 3 * 16)), - ( True, False, [4, 100, 3, 3, 16], (100, 4 * 3 * 3 * 16)), - (False, True, [10, 10, 10, 16], (1000, 16)), - (False, True, [10, 10, 10, 3, 16], (1000, 3 * 16)), - (False, True, [10, 10, 10, 3, 3, 16], (1000, 3 * 3 * 16)), - ( True, True, [4, 10, 10, 10, 16], (1000, 4 * 16)), - ( True, True, [4, 10, 10, 10, 3, 16], (1000, 4 * 3 * 16)), - ( True, True, [4, 10, 10, 10, 3, 3, 16], (1000, 4 * 3 * 3 * 16)), + (True, False, [4, 100, 16], (100, 4 * 16)), + (True, False, [4, 100, 3, 16], (100, 4 * 3 * 16)), + (True, False, [4, 100, 3, 3, 16], (100, 4 * 3 * 3 * 16)), + (False, True, [10, 10, 10, 16], (1000, 16)), + (False, True, [10, 10, 10, 3, 16], (1000, 3 * 16)), + (False, True, [10, 10, 10, 3, 3, 16], (1000, 3 * 3 * 16)), + (True, True, [4, 10, 10, 10, 16], (1000, 4 * 16)), + (True, True, [4, 10, 10, 10, 3, 16], (1000, 4 * 3 * 16)), + (True, True, [4, 10, 10, 10, 3, 3, 16], (1000, 4 * 3 * 3 * 16)), ], ) def test__to_vertexwise(is_time_series, is_voxel, size, desired_shape): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( - torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel + ) assert phlower_tensor.to_vertexwise()[0].shape == desired_shape @@ -171,31 +174,33 @@ def test__to_vertexwise(is_time_series, is_voxel, size, desired_shape): (False, False, [100, 16]), (False, False, [100, 3, 16]), (False, False, [100, 3, 3, 16]), - ( True, False, [4, 100, 16]), - ( True, False, [4, 100, 3, 16]), - ( True, False, [4, 100, 3, 3, 16]), - (False, True, [10, 10, 10, 16]), - (False, True, [10, 10, 10, 3, 16]), - (False, True, [10, 10, 10, 3, 3, 16]), - ( True, True, [4, 10, 10, 10, 16]), - ( True, True, [4, 10, 10, 10, 3, 16]), - ( True, True, [4, 10, 10, 10, 3, 3, 16]), + (True, False, [4, 100, 16]), + (True, False, [4, 100, 3, 16]), + (True, False, [4, 100, 3, 3, 16]), + (False, True, [10, 10, 10, 16]), + (False, True, [10, 10, 10, 3, 16]), + (False, True, [10, 10, 10, 3, 3, 16]), + (True, True, [4, 10, 10, 10, 16]), + (True, True, [4, 10, 10, 10, 3, 16]), + (True, True, [4, 10, 10, 10, 3, 3, 16]), ], ) def test__to_vertexwise_inverse(is_time_series, is_voxel, size): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( - torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel) - vertexwise, original_pattern, resultant_pattern, dict_shape\ - = phlower_tensor.to_vertexwise() + torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel + ) + vertexwise, original_pattern, resultant_pattern, dict_shape = ( + phlower_tensor.to_vertexwise() + ) assert len(vertexwise.shape) == 2 pattern = f"{resultant_pattern} -> {original_pattern}" actual = vertexwise.rearrange( - pattern, is_time_series=is_time_series, is_voxel=is_voxel, - **dict_shape) + pattern, is_time_series=is_time_series, is_voxel=is_voxel, **dict_shape + ) np.testing.assert_almost_equal( - actual.to_tensor().numpy(), - phlower_tensor.to_tensor().numpy()) + actual.to_tensor().numpy(), phlower_tensor.to_tensor().numpy() + ) @pytest.mark.parametrize( diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index f5ba46c..bcb8646 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -1,4 +1,3 @@ - import numpy as np import pytest import torch @@ -27,12 +26,16 @@ ) def test__spmm(size, is_time_series, repeat): phlower_tensor = PhlowerTensor( - torch.rand(*size), is_time_series=is_time_series) + torch.rand(*size), is_time_series=is_time_series + ) n = phlower_tensor.n_vertices() sparse = PhlowerTensor(torch.rand(n, n).to_sparse()) - actual_spmm = _functions.spmm( - sparse, phlower_tensor, repeat=repeat).to_tensor().numpy() + actual_spmm = ( + _functions.spmm(sparse, phlower_tensor, repeat=repeat) + .to_tensor() + .numpy() + ) sp_sparse = sp.coo_array(sparse.to_tensor().to_dense().numpy()) np_dense = phlower_tensor.to_tensor().numpy() diff --git a/tests/test_nn/test_core_modules/test_gcn.py b/tests/test_nn/test_core_modules/test_gcn.py index 53604a8..db767e4 100644 --- a/tests/test_nn/test_core_modules/test_gcn.py +++ b/tests/test_nn/test_core_modules/test_gcn.py @@ -7,7 +7,7 @@ def test__can_call_parameters(): - model = GCN(nodes=[4, 8], support_name='support') + model = GCN(nodes=[4, 8], support_name="support") # To check Concatenator inherit torch.nn.Module appropriately _ = model.parameters() @@ -26,14 +26,17 @@ def test__can_call_parameters(): ) def test__gcn(size, is_time_series): phlower_tensor = PhlowerTensor( - torch.rand(*size), is_time_series=is_time_series) - phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) + torch.rand(*size), is_time_series=is_time_series + ) + phlower_tensors = phlower_tensor_collection({"tensor": phlower_tensor}) n = phlower_tensor.n_vertices() - dict_supports = {'support': PhlowerTensor(torch.rand(n, n).to_sparse())} + dict_supports = {"support": PhlowerTensor(torch.rand(n, n).to_sparse())} model = GCN( nodes=[size[-1], size[-1], size[-1]], - support_name='support', activations=['tanh', 'identity']) + support_name="support", + activations=["tanh", "identity"], + ) actual = model(phlower_tensors, supports=dict_supports) diff --git a/tests/test_services/test_trainer/test_optimizer.py b/tests/test_services/test_trainer/test_optimizer.py new file mode 100644 index 0000000..a18122e --- /dev/null +++ b/tests/test_services/test_trainer/test_optimizer.py @@ -0,0 +1,121 @@ +from unittest import mock + +import pytest +import torch + +from phlower.services.trainer._optimizer import PhlowerOptimizerWrapper +from phlower.settings import PhlowerTrainerSetting + + +@pytest.mark.parametrize( + "optimizer, optimizer_parameters, schedulers", + [ + ("Adam", {"lr": 0.003}, []), + ("SGD", {"lr": 0.1, "weight_decay": 0.01}, []), + ( + "SGD", + {"lr": 0.1, "weight_decay": 0.01}, + [ + { + "scheduler": "ReduceLROnPlateau", + "parameters": {"mode": "min", "patience": 12}, + } + ], + ), + ], +) +def test__pass_kwargs_when_call_from_setting( + optimizer, optimizer_parameters, schedulers +): + setting = PhlowerTrainerSetting( + loss_setting={"name2loss": {}}, + optimizer_setting={ + "optimizer": optimizer, + "parameters": optimizer_parameters, + }, + scheduler_setting=schedulers, + ) + + with mock.patch.object( + PhlowerOptimizerWrapper, "__init__", return_value=None + ) as mocked: + model = torch.nn.Linear(in_features=10, out_features=10) + _ = PhlowerOptimizerWrapper.from_setting(setting, model=model) + + kwargs = mocked.call_args.kwargs + assert kwargs["optimizer"] == optimizer + assert kwargs["optimizer_kwargs"] == optimizer_parameters + + for scheduler in schedulers: + name = scheduler["scheduler"] + assert name in kwargs["schedulers"] + assert kwargs["schedulers"][name] == scheduler["parameters"] + + +@pytest.mark.parametrize( + "optimizer, lr, weight_decay, desired_optimizer", + [ + ("Adam", 0.001, 0, torch.optim.Adam), + ("Adam", 0.0003, 0, torch.optim.Adam), + ("SGD", 0.0005, 0.01, torch.optim.SGD), + ], +) +def test__optimizer_parameters(optimizer, lr, weight_decay, desired_optimizer): + model = torch.nn.Linear(in_features=10, out_features=10) + optimizer = PhlowerOptimizerWrapper( + parameters=model.parameters(), + optimizer=optimizer, + optimizer_kwargs={"lr": lr, "weight_decay": weight_decay}, + schedulers={}, + ) + + assert isinstance(optimizer._optimizer, desired_optimizer) + + state_dict = optimizer.state_dict() + assert len(state_dict["param_groups"]) == 1 + assert state_dict["param_groups"][0]["lr"] == lr + assert state_dict["param_groups"][0]["weight_decay"] == weight_decay + + +@pytest.mark.parametrize( + "schedulers, desired", + [ + ( + { + "ReduceLROnPlateau": {"mode": "min", "patience": 21}, + "StepLR": {"step_size": 30, "gamma": 0.2}, + }, + [ + torch.optim.lr_scheduler.ReduceLROnPlateau, + torch.optim.lr_scheduler.StepLR, + ], + ), + ( + { + "CosineAnnealingLR": {"T_max": 10}, + "ConstantLR": {"factor": 0.5}, + }, + [ + torch.optim.lr_scheduler.CosineAnnealingLR, + torch.optim.lr_scheduler.ConstantLR, + ], + ), + ], +) +def test__scheduler_parameters(schedulers, desired): + dummy = torch.nn.Linear(in_features=10, out_features=10) + optimizer = PhlowerOptimizerWrapper( + parameters=dummy.parameters(), + optimizer="Adam", + optimizer_kwargs={}, + schedulers=schedulers, + ) + + assert len(optimizer._schedulers) == len(schedulers) + for i, _scheduler in enumerate(desired): + isinstance(optimizer._schedulers[i], _scheduler) + + for i, _scheduler in enumerate(desired): + params = schedulers[_scheduler.__name__] + for k, v in params.items(): + assert getattr(optimizer._schedulers[i], k) == v From 4a76e63f27e667acd1badd9a7bb67b73d64bb2b7 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 13 Aug 2024 16:40:20 +0900 Subject: [PATCH 34/89] add format check in linting --- Makefile | 1 + 1 file changed, 1 insertion(+) diff --git a/Makefile b/Makefile index 376f6fe..e1d8a55 100644 --- a/Makefile +++ b/Makefile @@ -30,6 +30,7 @@ e2e_test: .PHONY: lint lint: poetry run python3 -m ruff check --diff + poetry run python3 -m ruff format --diff # $(MAKE) mypy .PHONY: dev-install From 6e47dc9744f0a67c89766ca4b79a4e6afe9fa21a Mon Sep 17 00:00:00 2001 From: horiem Date: Tue, 13 Aug 2024 16:40:51 +0900 Subject: [PATCH 35/89] update test_spmm --- tests/test_nn/test_core_modules/test_functions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index f5ba46c..f7de675 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -42,7 +42,9 @@ def assert_correct(actual, array): desired = array for _ in range(repeat): desired = sp_sparse @ desired - np.testing.assert_almost_equal(actual, desired, decimal=5) + norm = np.mean(np.linalg.norm(desired, axis=-1)) + np.testing.assert_almost_equal( + actual / norm, desired / norm, decimal=5) return for i in range(array.shape[1]): From d5cfa620fead8a819e4f98685ee22f2d6eeeaa99 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 13 Aug 2024 19:42:35 +0900 Subject: [PATCH 36/89] fix lint warnings --- tests/test_nn/test_core_modules/test_functions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 996e986..4f355ce 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -47,7 +47,8 @@ def assert_correct(actual, array): desired = sp_sparse @ desired norm = np.mean(np.linalg.norm(desired, axis=-1)) np.testing.assert_almost_equal( - actual / norm, desired / norm, decimal=5) + actual / norm, desired / norm, decimal=5 + ) return for i in range(array.shape[1]): From 938439f14ba1f6a4db6653b2283476ed84621c85 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 14 Aug 2024 12:46:48 +0900 Subject: [PATCH 37/89] add tests for module settings --- src/phlower/settings/_interface.py | 46 +++++- .../_module_settings/_concatenator_setting.py | 5 +- .../settings/_module_settings/_gcn_setting.py | 27 ++- .../settings/_module_settings/_mlp_setting.py | 38 ++++- .../check_concatenator_nodes.yml | 63 +++++++ .../data/gcn_setting/check_gcn_nodes.yml | 54 ++++++ .../data/mlp_setting/check_mlp_nodes.yml | 50 ++++++ .../check_gcn_share_nodes.yml | 0 .../check_mlp_share_nodes.yml | 0 .../with_share_nn.yml | 0 .../test_concatenator_setting.py | 84 ++++++++++ .../test_module_settings/test_gcn_settings.py | 156 ++++++++++++++++++ .../test_module_settings/test_mlp_setting.py | 135 +++++++++++++++ .../test_share_setting.py | 72 ++++++++ .../test_share_settings.py | 30 ---- 15 files changed, 711 insertions(+), 49 deletions(-) create mode 100644 tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml create mode 100644 tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml create mode 100644 tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml rename tests/test_settings/test_module_settings/data/{share_settings => share_setting}/check_gcn_share_nodes.yml (100%) rename tests/test_settings/test_module_settings/data/{share_settings => share_setting}/check_mlp_share_nodes.yml (100%) rename tests/test_settings/test_module_settings/data/{share_settings => share_setting}/with_share_nn.yml (100%) create mode 100644 tests/test_settings/test_module_settings/test_concatenator_setting.py create mode 100644 tests/test_settings/test_module_settings/test_gcn_settings.py create mode 100644 tests/test_settings/test_module_settings/test_mlp_setting.py create mode 100644 tests/test_settings/test_module_settings/test_share_setting.py delete mode 100644 tests/test_settings/test_module_settings/test_share_settings.py diff --git a/src/phlower/settings/_interface.py b/src/phlower/settings/_interface.py index 23acbe8..ac11de8 100644 --- a/src/phlower/settings/_interface.py +++ b/src/phlower/settings/_interface.py @@ -28,17 +28,53 @@ def resolve( class IPhlowerLayerParameters(metaclass=abc.ABCMeta): @abc.abstractmethod - def gather_input_dims(self, *input_dims: int) -> int: ... + def gather_input_dims(self, *input_dims: int) -> int: + """Gather feature dimensions of input tensors + + Args: + *input_dims (list[int]): feature dimensions of input tensors + + Returns: + int: resolved dimension of input tensors + """ + ... @abc.abstractmethod - def get_n_nodes(self) -> list[int] | None: ... + def get_n_nodes(self) -> list[int] | None: + """Return feature dimensions inside PhlowerLayer. + + Returns: + list[int] | None: feature dimensions + """ + ... @abc.abstractmethod - def overwrite_nodes(self, nodes: list[int]) -> None: ... + def overwrite_nodes(self, nodes: list[int]) -> None: + """overwrite feature dimensions by using resolved information + + Args: + nodes (list[int]): + feature dimensions which is resolved + by using status of precedent and succedent phlower modules. + """ + ... @property @abc.abstractmethod - def need_reference(self) -> bool: ... + def need_reference(self) -> bool: + """Whether the reference to other phlower modules is needed or not. + + Returns: + bool: True if reference to other phlower modules is necessary. + """ + ... @abc.abstractmethod - def get_reference(self, parent: IReadOnlyReferenceGroupSetting): ... + def get_reference(self, parent: IReadOnlyReferenceGroupSetting): + """Get reference information from parent group setting + + Args: + parent (IReadOnlyReferenceGroupSetting): + Reference to group setting of its parent + """ + ... diff --git a/src/phlower/settings/_module_settings/_concatenator_setting.py b/src/phlower/settings/_module_settings/_concatenator_setting.py index 341e2a1..ed05309 100644 --- a/src/phlower/settings/_module_settings/_concatenator_setting.py +++ b/src/phlower/settings/_module_settings/_concatenator_setting.py @@ -10,9 +10,8 @@ class ConcatenatorSetting(IPhlowerLayerParameters, pydantic.BaseModel): - nodes: list[int] | None = Field( - None - ) # This property only overwritten when resolving. + # This property only overwritten when resolving. + nodes: list[int] | None = Field(None) activation: str = Field("identity", frozen=True) # special keyward to forbid extra fields in pydantic diff --git a/src/phlower/settings/_module_settings/_gcn_setting.py b/src/phlower/settings/_module_settings/_gcn_setting.py index 6bb204a..4245771 100644 --- a/src/phlower/settings/_module_settings/_gcn_setting.py +++ b/src/phlower/settings/_module_settings/_gcn_setting.py @@ -11,9 +11,8 @@ class GCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): - nodes: list[int] = Field( - ... - ) # This property only overwritten when resolving. + # This property only overwritten when resolving. + nodes: list[int] = Field(...) support_name: str = Field(..., frozen=True) repeat: int = Field(1, frozen=True) factor: float = Field(1.0, frozen=True) @@ -52,6 +51,21 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: return vals + @pydantic.model_validator(mode="before") + @classmethod + def fill_empty_activations_dropouts(cls, values: dict): + n_nodes = len(values.get("nodes")) + activations = values.get("activations", []) + dropouts = values.get("dropouts", []) + + if len(activations) == 0: + values["activations"] = ["identity" for _ in range(n_nodes - 1)] + + if len(dropouts) == 0: + values["dropouts"] = [0 for _ in range(n_nodes - 1)] + + return values + @pydantic.model_validator(mode="after") def check_nodes_size(self) -> Self: if len(self.nodes) - 1 != len(self.activations): @@ -60,6 +74,13 @@ def check_nodes_size(self) -> Self: "in GCNSettings." " len(nodes) must be equal to 1 + len(activations)." ) + + if len(self.nodes) - 1 != len(self.dropouts): + raise ValueError( + "Size of nodes and dropouts is not compatible " + "in GCNSettings." + " len(nodes) must be equal to 1 + len(dropouts)." + ) return self def get_n_nodes(self) -> list[int]: diff --git a/src/phlower/settings/_module_settings/_mlp_setting.py b/src/phlower/settings/_module_settings/_mlp_setting.py index b4d02ce..1c168eb 100644 --- a/src/phlower/settings/_module_settings/_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_mlp_setting.py @@ -11,24 +11,23 @@ class MLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): - nodes: list[int] = Field( - ... - ) # This property only overwritten when resolving. + # This property only overwritten when resolving. + nodes: list[int] = Field(...) activations: list[str] = Field(default_factory=lambda: [], frozen=True) dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) bias: bool = Field(False, frozen=True) def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: - raise ValueError("only one input is allowed in GCN.") + raise ValueError("only one input is allowed in MLP.") return input_dims[0] - @pydantic.field_validator("nodes") + @pydantic.field_validator("nodes", mode="before") @classmethod def check_n_nodes(cls, vals: list[int]) -> list[int]: if len(vals) < 2: raise ValueError( - "size of nodes must be larger than 1 in GCNSettings." + "size of nodes must be larger than 1 in MLPSettings." f" input: {vals}" ) @@ -40,20 +39,43 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: continue raise ValueError( - "nodes in GCN is inconsistent. " + "nodes in MLP is inconsistent. " f"value {v} in {i}-th of nodes is not allowed." ) return vals + @pydantic.model_validator(mode="before") + @classmethod + def fill_empty_activations_dropouts(cls, values: dict): + n_nodes = len(values.get("nodes")) + activations = values.get("activations", []) + dropouts = values.get("dropouts", []) + + if len(activations) == 0: + values["activations"] = ["identity" for _ in range(n_nodes - 1)] + + if len(dropouts) == 0: + values["dropouts"] = [0 for _ in range(n_nodes - 1)] + + return values + @pydantic.model_validator(mode="after") def check_nodes_size(self) -> Self: if len(self.nodes) - 1 != len(self.activations): raise ValueError( "Size of nodes and activations is not compatible " - "in GCNSettings." + "in MLPSettings." " len(nodes) must be equal to 1 + len(activations)." ) + + if len(self.nodes) - 1 != len(self.dropouts): + raise ValueError( + "Size of nodes and dropouts is not compatible " + "in MLPSettings." + " len(nodes) must be equal to 1 + len(dropouts)." + ) + return self def get_n_nodes(self) -> list[int] | None: diff --git a/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml b/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml new file mode 100644 index 0000000..50c0c45 --- /dev/null +++ b/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml @@ -0,0 +1,63 @@ +misc: + + tests: + MLP0: [10, 20, 15] + MLP1: [12, 20, 25] + Concat0: [40, 40] + MLP2: [40, 20, 5] + + +model: + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - Concat0 + nn_parameters: + nodes: [-1, 20, 15] + activations: ["Identity", "identity"] + + - nn_type: MLP + name: MLP1 + input_keys: + - feature1 + output_key: mlp1 + destinations: + - Concat0 + nn_parameters: + nodes: [-1, 20, 25] + activations: ["Identity", "identity"] + + - nn_type: Concatenator + name: Concat0 + input_keys: + - mlp0 + - mlp1 + output_key: concat0 + destinations: + - MLP2 + + - nn_type: MLP + name: MLP2 + input_keys: + - concat0 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 20, 5] + activations: ["Identity", "identity"] diff --git a/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml b/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml new file mode 100644 index 0000000..d04502d --- /dev/null +++ b/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml @@ -0,0 +1,54 @@ +misc: + + tests: + GCN0: [12, 20, 5] + GCN1: [5, 20, 5] + GCN2: [5, 5] + +model: + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: GCN + name: GCN0 + input_keys: + - feature1 + output_key: gcn0 + destinations: + - GCN1 + nn_parameters: + nodes: [-1, 20, 5] + activations: ["Identity", "identity"] + support_name: support1 + + - nn_type: GCN + name: GCN1 + input_keys: + - gcn0 + output_key: gcn1 + destinations: + - GCN2 + nn_parameters: + nodes: [-1, 20, 5] + support_name: support1 + + - nn_type: GCN + name: GCN2 + input_keys: + - gcn1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + diff --git a/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml b/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml new file mode 100644 index 0000000..26d5c38 --- /dev/null +++ b/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml @@ -0,0 +1,50 @@ +misc: + + tests: + MLP0: [12, 20, 10] + MLP1: [10, 20, 30] + MLP2: [30, 20, 5] + +model: + network: + nn_type: GROUP + name: SAMPLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature1 + output_key: mlp0 + destinations: + - MLP1 + nn_parameters: + nodes: [-1, 20, 10] + activations: ["Identity", "identity"] + + - nn_type: MLP + name: MLP1 + input_keys: + - mlp0 + output_key: mlp1 + destinations: + - MLP2 + nn_parameters: + nodes: [-1, 20, 30] + + - nn_type: MLP + name: MLP2 + input_keys: + - mlp1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 20, 5] diff --git a/tests/test_settings/test_module_settings/data/share_settings/check_gcn_share_nodes.yml b/tests/test_settings/test_module_settings/data/share_setting/check_gcn_share_nodes.yml similarity index 100% rename from tests/test_settings/test_module_settings/data/share_settings/check_gcn_share_nodes.yml rename to tests/test_settings/test_module_settings/data/share_setting/check_gcn_share_nodes.yml diff --git a/tests/test_settings/test_module_settings/data/share_settings/check_mlp_share_nodes.yml b/tests/test_settings/test_module_settings/data/share_setting/check_mlp_share_nodes.yml similarity index 100% rename from tests/test_settings/test_module_settings/data/share_settings/check_mlp_share_nodes.yml rename to tests/test_settings/test_module_settings/data/share_setting/check_mlp_share_nodes.yml diff --git a/tests/test_settings/test_module_settings/data/share_settings/with_share_nn.yml b/tests/test_settings/test_module_settings/data/share_setting/with_share_nn.yml similarity index 100% rename from tests/test_settings/test_module_settings/data/share_settings/with_share_nn.yml rename to tests/test_settings/test_module_settings/data/share_setting/with_share_nn.yml diff --git a/tests/test_settings/test_module_settings/test_concatenator_setting.py b/tests/test_settings/test_module_settings/test_concatenator_setting.py new file mode 100644 index 0000000..5034252 --- /dev/null +++ b/tests/test_settings/test_module_settings/test_concatenator_setting.py @@ -0,0 +1,84 @@ +import pathlib + +import hypothesis.strategies as st +import pytest +import yaml +from hypothesis import assume, given, settings + +from phlower.settings import PhlowerModelSetting +from phlower.settings._module_settings import ConcatenatorSetting + + +@pytest.mark.parametrize("nodes", [(None), ([10, 10])]) +def test__can_accept_valid_n_nodes(nodes): + _ = ConcatenatorSetting(nodes=nodes) + + +@pytest.mark.parametrize("nodes", [([5]), ([10, 10, 10])]) +def test__raise_error_when_invalid_n_nodes(nodes): + with pytest.raises(ValueError): + _ = ConcatenatorSetting(nodes=nodes) + + +@pytest.mark.parametrize( + "input_dims, desired", [([30, 50, 40], 120), ([40], 40), ([100, 10], 110)] +) +def test__gather_input_dims(input_dims, desired): + setting = ConcatenatorSetting() + + assert setting.gather_input_dims(*input_dims) == desired + + +@st.composite +def same_length_lists(draw): + n_elements = draw(st.integers(min_value=2, max_value=2)) + fixed_length_list = st.lists( + st.integers(min_value=1, max_value=200), + min_size=n_elements, + max_size=n_elements, + ) + + return (draw(fixed_length_list), draw(fixed_length_list)) + + +@given(same_length_lists()) +@settings(max_examples=100) +def test__nodes_is_update_after_overwrite_nodes(lists): + nodes, update_nodes = lists + assume(nodes != update_nodes) + setting = ConcatenatorSetting() + + before_nodes = setting.get_n_nodes() + assert before_nodes != update_nodes + + setting.overwrite_nodes(update_nodes) + assert setting.get_n_nodes() == update_nodes + + +def test__reference_is_not_necessary(): + setting = ConcatenatorSetting() + + assert not setting.need_reference + + +# region E2E tests only for MLPSettings + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/concatenator_setting" + + +@pytest.mark.parametrize("yaml_file", ["check_concatenator_nodes.yml"]) +def test__nodes_after_resolve(yaml_file): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + + assert len(content["misc"]["tests"].items()) > 0 + + for key, value in content["misc"]["tests"].items(): + target = setting.network.search_module_setting(key) + assert target.get_n_nodes() == value + + +# endregion diff --git a/tests/test_settings/test_module_settings/test_gcn_settings.py b/tests/test_settings/test_module_settings/test_gcn_settings.py new file mode 100644 index 0000000..0061f57 --- /dev/null +++ b/tests/test_settings/test_module_settings/test_gcn_settings.py @@ -0,0 +1,156 @@ +import pathlib + +import hypothesis.strategies as st +import pytest +import yaml +from hypothesis import assume, given, settings + +from phlower.settings import PhlowerModelSetting +from phlower.settings._module_settings import GCNSetting + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), + ([10, 30], ["identity"], [0.3]), + ([5, 10, 20, 5], ["relu", "relu", "tanh"], [0.3, 0.2, 0.1]), + ([-1, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), + ], +) +def test__can_accept_valid_n_nodes(nodes, activations, dropouts): + _ = GCNSetting( + nodes=nodes, + support_name="dummy", + activations=activations, + dropouts=dropouts, + ) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["identity"], []), + ([10, 30], [], [0.3, 0.4]), + ([5, 10, 20, 5], ["relu", "relu", "tanh", "identity"], [0.3, 0.2, 0.1]), + ([5, -1, 20, 5], ["relu", "relu", "tanh"], [0.3, 0.2, 0.1]), + ([5, 10, 20, 5], ["relu", "relu", "tanh"], [0.3]), + ([10], [], []), + ], +) +def test__raise_error_when_invalid_n_nodes(nodes, activations, dropouts): + with pytest.raises(ValueError): + _ = GCNSetting( + nodes=nodes, + support_name="dummy", + activations=activations, + dropouts=dropouts, + ) + + +@pytest.mark.parametrize("input_dims", [([30]), ([40]), ([100])]) +def test__gather_input_dims(input_dims): + setting = GCNSetting( + nodes=[10, 20], + support_name="dummy", + activations=["identity"], + dropouts=[0.1], + ) + + assert setting.gather_input_dims(*input_dims) == input_dims[0] + + +@pytest.mark.parametrize("input_dims", [([]), ([40, 400]), ([10, 0, 1])]) +def test__raise_error_invalid_input_dims(input_dims): + setting = GCNSetting( + nodes=[10, 20], + support_name="dummy", + activations=["identity"], + dropouts=[0.1], + ) + + with pytest.raises(ValueError): + _ = setting.gather_input_dims(*input_dims) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [([10, 20, 30], [], []), ([10, 20, 30, 40, 50], [], [])], +) +def test__fill_default_settings(nodes, activations, dropouts): + setting = GCNSetting( + nodes=nodes, + support_name="dummy", + activations=activations, + dropouts=dropouts, + ) + desired_activations = ["identity" for _ in range(len(nodes) - 1)] + desired_dropouts = [0 for _ in range(len(nodes) - 1)] + + assert setting.activations == desired_activations + assert setting.dropouts == desired_dropouts + + +@st.composite +def same_length_lists(draw): + n_elements = draw(st.integers(min_value=2, max_value=10)) + fixed_length_list = st.lists( + st.integers(min_value=1, max_value=200), + min_size=n_elements, + max_size=n_elements, + ) + + return (draw(fixed_length_list), draw(fixed_length_list)) + + +@given(same_length_lists()) +@settings(max_examples=100) +def test__nodes_is_update_after_overwrite_nodes(lists): + nodes, update_nodes = lists + assume(nodes != update_nodes) + setting = GCNSetting( + nodes=nodes, + support_name="dummy", + activations=["identity" for _ in range(len(nodes) - 1)], + dropouts=[0.1 for _ in range(len(nodes) - 1)], + ) + + before_nodes = setting.get_n_nodes() + assert before_nodes != update_nodes + + setting.overwrite_nodes(update_nodes) + assert setting.get_n_nodes() == update_nodes + + +def test__reference_is_not_necessary(): + setting = GCNSetting( + nodes=[10, 20], + support_name="dummy", + activations=["identity"], + dropouts=[0.1], + ) + + assert not setting.need_reference + + +# region E2E tests only for MLPSettings + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/gcn_setting" + + +@pytest.mark.parametrize("yaml_file", ["check_gcn_nodes.yml"]) +def test__nodes_after_resolve(yaml_file): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + + assert len(content["misc"]["tests"].items()) > 0 + + for key, value in content["misc"]["tests"].items(): + target = setting.network.search_module_setting(key) + assert target.get_n_nodes() == value + + +# endregion diff --git a/tests/test_settings/test_module_settings/test_mlp_setting.py b/tests/test_settings/test_module_settings/test_mlp_setting.py new file mode 100644 index 0000000..0835a2f --- /dev/null +++ b/tests/test_settings/test_module_settings/test_mlp_setting.py @@ -0,0 +1,135 @@ +import pathlib + +import hypothesis.strategies as st +import pytest +import yaml +from hypothesis import assume, given, settings + +from phlower.settings import PhlowerModelSetting +from phlower.settings._module_settings import MLPSetting + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/share_settings" + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), + ([10, 30], ["identity"], [0.3]), + ([5, 10, 20, 5], ["relu", "relu", "tanh"], [0.3, 0.2, 0.1]), + ([-1, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), + ], +) +def test__can_accept_valid_n_nodes(nodes, activations, dropouts): + _ = MLPSetting(nodes=nodes, activations=activations, dropouts=dropouts) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["identity"], []), + ([10, 30], [], [0.3, 0.4]), + ([5, 10, 20, 5], ["relu", "relu", "tanh", "identity"], [0.3, 0.2, 0.1]), + ([5, -1, 20, 5], ["relu", "relu", "tanh"], [0.3, 0.2, 0.1]), + ([5, 10, 20, 5], ["relu", "relu", "tanh"], [0.3]), + ([10], [], []), + ], +) +def test__raise_error_when_invalid_n_nodes(nodes, activations, dropouts): + with pytest.raises(ValueError): + _ = MLPSetting(nodes=nodes, activations=activations, dropouts=dropouts) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [([10, 20, 30], [], []), ([10, 20, 30, 40, 50], [], [])], +) +def test__fill_default_settings(nodes, activations, dropouts): + setting = MLPSetting( + nodes=nodes, activations=activations, dropouts=dropouts + ) + desired_activations = ["identity" for _ in range(len(nodes) - 1)] + desired_dropouts = [0 for _ in range(len(nodes) - 1)] + + assert setting.activations == desired_activations + assert setting.dropouts == desired_dropouts + + +@pytest.mark.parametrize("input_dims", [([30]), ([40]), ([100])]) +def test__gather_input_dims(input_dims): + setting = MLPSetting( + nodes=[10, 20], activations=["identity"], dropouts=[0.1] + ) + + assert setting.gather_input_dims(*input_dims) == input_dims[0] + + +@pytest.mark.parametrize("input_dims", [([]), ([40, 400]), ([10, 0, 1])]) +def test__raise_error_invalid_input_dims(input_dims): + setting = MLPSetting( + nodes=[10, 20], activations=["identity"], dropouts=[0.1] + ) + + with pytest.raises(ValueError): + _ = setting.gather_input_dims(*input_dims) + + +@st.composite +def same_length_lists(draw): + n_elements = draw(st.integers(min_value=2, max_value=10)) + fixed_length_list = st.lists( + st.integers(min_value=1, max_value=200), + min_size=n_elements, + max_size=n_elements, + ) + + return (draw(fixed_length_list), draw(fixed_length_list)) + + +@given(same_length_lists()) +@settings(max_examples=100) +def test__nodes_is_update_after_overwrite_nodes(lists): + nodes, update_nodes = lists + assume(nodes != update_nodes) + setting = MLPSetting( + nodes=nodes, + activations=["identity" for _ in range(len(nodes) - 1)], + dropouts=[0.1 for _ in range(len(nodes) - 1)], + ) + + before_nodes = setting.get_n_nodes() + assert before_nodes != update_nodes + + setting.overwrite_nodes(update_nodes) + assert setting.get_n_nodes() == update_nodes + + +def test__reference_is_not_necessary(): + setting = MLPSetting( + nodes=[10, 20], activations=["identity"], dropouts=[0.1] + ) + + assert not setting.need_reference + + +# region E2E tests only for MLPSettings + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/mlp_setting" + + +@pytest.mark.parametrize("yaml_file", ["check_mlp_nodes.yml"]) +def test__nodes_after_resolve(yaml_file): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + + assert len(content["misc"]["tests"].items()) > 0 + + for key, value in content["misc"]["tests"].items(): + target = setting.network.search_module_setting(key) + assert target.get_n_nodes() == value + + +# endregion diff --git a/tests/test_settings/test_module_settings/test_share_setting.py b/tests/test_settings/test_module_settings/test_share_setting.py new file mode 100644 index 0000000..d842efa --- /dev/null +++ b/tests/test_settings/test_module_settings/test_share_setting.py @@ -0,0 +1,72 @@ +import pathlib +from unittest import mock + +import hypothesis.strategies as st +import pytest +import yaml +from hypothesis import given, settings + +from phlower.settings import PhlowerModelSetting, PhlowerSetting +from phlower.settings._module_settings import ShareSetting + + +@given(st.lists(st.integers(), max_size=100)) +@settings(max_examples=100) +def test__gather_input_dims_of_reference_setting(input_dims): + setting = ShareSetting(reference_name="dummy") + mocked = mock.MagicMock() + setting.reference = mocked + + _ = setting.gather_input_dims(*input_dims) + + mocked.gather_input_dims.assert_called_once_with(*input_dims) + + +def test__get_n_nodes_of_reference_setting(): + setting = ShareSetting(reference_name="dummy") + mocked = mock.MagicMock() + setting.reference = mocked + + _ = setting.get_n_nodes() + + mocked.get_n_nodes.assert_called_once() + + +@pytest.mark.parametrize("reference_name", ["dummy", "mlp0", "gcn0"]) +def test__call_parent_function_when_get_reference(reference_name): + setting = ShareSetting(reference_name=reference_name) + mocked = mock.MagicMock() + + setting.get_reference(mocked) + + mocked.search_module_setting.assert_called_once_with(reference_name) + + +# region E2E tests only for ShareSettings + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/share_setting" + + +def test__can_resolve(): + setting = PhlowerSetting.read_yaml(_TEST_DATA_DIR / "with_share_nn.yml") + setting.model.network.resolve(is_first=True) + + +@pytest.mark.parametrize( + "yaml_file", ["check_gcn_share_nodes.yml", "check_mlp_share_nodes.yml"] +) +def test__nodes_after_resolve(yaml_file): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + + assert len(content["misc"]["tests"].items()) > 0 + + for key, value in content["misc"]["tests"].items(): + target = setting.network.search_module_setting(key) + assert target.get_n_nodes() == value + + +# endregion diff --git a/tests/test_settings/test_module_settings/test_share_settings.py b/tests/test_settings/test_module_settings/test_share_settings.py deleted file mode 100644 index 9355a13..0000000 --- a/tests/test_settings/test_module_settings/test_share_settings.py +++ /dev/null @@ -1,30 +0,0 @@ -import pathlib - -import pytest -import yaml - -from phlower.settings import PhlowerModelSetting, PhlowerSetting - -_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/share_settings" - - -def test__can_resolve(): - setting = PhlowerSetting.read_yaml(_TEST_DATA_DIR / "with_share_nn.yml") - setting.model.network.resolve(is_first=True) - - -@pytest.mark.parametrize( - "yaml_file", ["check_gcn_share_nodes.yml", "check_mlp_share_nodes.yml"] -) -def test__nodes_after_resolve(yaml_file): - with open(_TEST_DATA_DIR / yaml_file) as fr: - content = yaml.load(fr, Loader=yaml.SafeLoader) - - setting = PhlowerModelSetting(**content["model"]) - setting.network.resolve(is_first=True) - - assert len(content["misc"]["tests"].items()) > 0 - - for key, value in content["misc"]["tests"].items(): - target = setting.network.search_module_setting(key) - assert target.get_n_nodes() == value From 475d7862ea586eab457c6581e102f9fae3bbd5f9 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 14 Aug 2024 17:12:30 +0900 Subject: [PATCH 38/89] add SimulationField to manage all constant field values. --- src/phlower/__init__.py | 1 + src/phlower/_fields/__init__.py | 1 + src/phlower/_fields/_simulation_field.py | 46 ++++++++++++ src/phlower/data/_collate_fn.py | 16 ++--- src/phlower/data/_datasets.py | 13 ++-- src/phlower/data/_lumped_data.py | 15 ++-- src/phlower/nn/_core_modules/_concatenator.py | 3 +- src/phlower/nn/_core_modules/_gcn.py | 5 +- src/phlower/nn/_core_modules/_mlp.py | 2 + src/phlower/nn/_core_modules/_share.py | 5 +- src/phlower/nn/_group_module.py | 6 +- src/phlower/nn/_interface_module.py | 3 +- src/phlower/nn/_phlower_module_adpter.py | 8 +-- src/phlower/services/predictor/_predictor.py | 4 +- src/phlower/services/trainer/_trainer.py | 2 +- src/phlower/settings/__init__.py | 2 +- src/phlower/settings/_model_setting.py | 71 +++++++++++++++++++ src/phlower/settings/_phlower_setting.py | 20 +----- tests/test_data/test_data_loader.py | 16 ++--- tests/test_data/test_datasets.py | 18 ++--- tests/test_nn/test_core_modules/test_gcn.py | 2 +- tests/test_nn/test_group_module.py | 4 +- .../simple_module_with_simultion_field.yml | 52 ++++++++++++++ 23 files changed, 236 insertions(+), 79 deletions(-) create mode 100644 src/phlower/_fields/__init__.py create mode 100644 src/phlower/_fields/_simulation_field.py create mode 100644 src/phlower/settings/_model_setting.py create mode 100644 tests/test_settings/data/groups/simple_module_with_simultion_field.yml diff --git a/src/phlower/__init__.py b/src/phlower/__init__.py index 987a5c3..1dbbff2 100644 --- a/src/phlower/__init__.py +++ b/src/phlower/__init__.py @@ -6,6 +6,7 @@ phlower_tensor, ) from phlower._base.array import IPhlowerArray +from phlower._fields import ISimulationField from phlower.version import __version__ __all__ = ["__version__"] diff --git a/src/phlower/_fields/__init__.py b/src/phlower/_fields/__init__.py new file mode 100644 index 0000000..1fe9b26 --- /dev/null +++ b/src/phlower/_fields/__init__.py @@ -0,0 +1 @@ +from phlower._fields._simulation_field import ISimulationField, SimulationField diff --git a/src/phlower/_fields/_simulation_field.py b/src/phlower/_fields/_simulation_field.py new file mode 100644 index 0000000..a572a3c --- /dev/null +++ b/src/phlower/_fields/_simulation_field.py @@ -0,0 +1,46 @@ +import abc + +from phlower import PhlowerTensor +from phlower._base import GraphBatchInfo +from phlower.collections.tensors import IPhlowerTensorCollections + + +class ISimulationField(metaclass=abc.ABCMeta): + @abc.abstractmethod + def __getitem__(self, name: str) -> PhlowerTensor: ... + + @abc.abstractmethod + def keys(self): ... + + +class SimulationField(ISimulationField): + def __init__( + self, + field_tensors: IPhlowerTensorCollections, + batch_info: dict[str, GraphBatchInfo] | None = None, + ) -> None: + self._field_tensors = field_tensors + + if batch_info is None: + batch_info = {} + self._batch_info = batch_info + + def keys(self): + return self._field_tensors.keys() + + def __getitem__(self, name: str) -> PhlowerTensor: + if name not in self._field_tensors: + raise KeyError(f"{name} is not found in simulation field.") + return self._field_tensors[name] + + def get_batch_info(self, name: str) -> GraphBatchInfo: + if name not in self._field_tensors: + raise KeyError(f"{name} is not found in simulation field.") + return self._batch_info[name] + + # HACK: Under construction + # def calculate_laplacians(self, target: PhlowerTensor): + # ... + + # def calculate_gradient(self, target: PhlowerTensor): + # ... diff --git a/src/phlower/data/_collate_fn.py b/src/phlower/data/_collate_fn.py index 87d359a..1c8148a 100644 --- a/src/phlower/data/_collate_fn.py +++ b/src/phlower/data/_collate_fn.py @@ -19,7 +19,7 @@ def __init__( def __call__(self, batch: list[LumpedArrayData]) -> LumpedTensorData: inputs = SequencedDictArray([v.x_data for v in batch]) outputs = SequencedDictArray([v.y_data for v in batch]) - sparse_supports = SequencedDictArray([v.sparse_supports for v in batch]) + field_data = SequencedDictArray([v.field_data for v in batch]) # concatenate and send inputs_tensors, inputs_batch_info = inputs.to_batched_tensor( @@ -33,21 +33,19 @@ def __call__(self, batch: list[LumpedArrayData]) -> LumpedTensorData: non_blocking=self._non_blocking, dimensions=self._dimensions, ) - sparse_supports_tensors, support_batch_info = ( - sparse_supports.to_batched_tensor( - device=self._device, - non_blocking=self._non_blocking, - dimensions=self._dimensions, - ) + field_tensors, field_batch_info = field_data.to_batched_tensor( + device=self._device, + non_blocking=self._non_blocking, + dimensions=self._dimensions, ) data_directories = [b.data_directory for b in batch] return LumpedTensorData( x_data=inputs_tensors, y_data=outputs_tensors, - sparse_supports=sparse_supports_tensors, + field_data=field_tensors, data_directories=data_directories, x_batch_info=inputs_batch_info, y_batch_info=outputs_batch_info, - supports_batch_info=support_batch_info, + field_batch_info=field_batch_info, ) diff --git a/src/phlower/data/_datasets.py b/src/phlower/data/_datasets.py index 792e49f..8a0cdad 100644 --- a/src/phlower/data/_datasets.py +++ b/src/phlower/data/_datasets.py @@ -23,17 +23,16 @@ def __init__( y_variable_names: list[str] | None, directories: list[pathlib.Path], *, - support_names: list[str] = None, + field_names: list[str] = None, allow_no_y_data: bool = False, - decrypt_key: bytes | None = None, - **kwargs, + decrypt_key: bytes | None = None ): self._x_variable_names = x_variable_names if y_variable_names is None: y_variable_names = [] self._y_varaible_names = y_variable_names self._directories = [PhlowerDirectory(d) for d in directories] - self._support_names = support_names if support_names is not None else [] + self._field_names = field_names if field_names is not None else [] self._allow_no_y_data = allow_no_y_data self._decrypt_key = decrypt_key @@ -51,13 +50,13 @@ def __getitem__(self, idx: int) -> LumpedArrayData: self._y_varaible_names, allow_missing=self._allow_no_y_data, ) - support_data = self._load_data( - data_directory, self._support_names, allow_missing=False + field_data = self._load_data( + data_directory, self._field_names, allow_missing=False ) return LumpedArrayData( x_data=x_data, y_data=y_data, - sparse_supports=support_data, + field_data=field_data, data_directory=data_directory, ) diff --git a/src/phlower/data/_lumped_data.py b/src/phlower/data/_lumped_data.py index e310815..b67eeb3 100644 --- a/src/phlower/data/_lumped_data.py +++ b/src/phlower/data/_lumped_data.py @@ -1,5 +1,6 @@ from phlower import IPhlowerArray from phlower._base import GraphBatchInfo +from phlower._fields import SimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.io import PhlowerDirectory @@ -9,12 +10,12 @@ def __init__( self, x_data: dict[str, IPhlowerArray], y_data: dict[str, IPhlowerArray], - sparse_supports: dict[str, IPhlowerArray], + field_data: dict[str, IPhlowerArray], data_directory: PhlowerDirectory, ) -> None: self.x_data = x_data self.y_data = y_data - self.sparse_supports = sparse_supports + self.field_data = field_data self.data_directory = data_directory @@ -22,18 +23,20 @@ class LumpedTensorData: def __init__( self, x_data: IPhlowerTensorCollections, - sparse_supports: IPhlowerTensorCollections, + field_data: IPhlowerTensorCollections, data_directories: list[PhlowerDirectory] | None = None, y_data: IPhlowerTensorCollections | None = None, x_batch_info: dict[str, GraphBatchInfo] | None = None, y_batch_info: dict[str, GraphBatchInfo] | None = None, - supports_batch_info: dict[str, GraphBatchInfo] | None = None, + field_batch_info: dict[str, GraphBatchInfo] | None = None, ) -> None: self.x_data = x_data self.y_data = y_data - self.sparse_supports = sparse_supports self.data_directories = data_directories self.x_batch_info = x_batch_info self.y_batch_info = y_batch_info - self.supports_batch_info = supports_batch_info + + self.field_data = SimulationField( + field_tensors=field_data, batch_info=field_batch_info + ) diff --git a/src/phlower/nn/_core_modules/_concatenator.py b/src/phlower/nn/_core_modules/_concatenator.py index 4fc600b..d735f14 100644 --- a/src/phlower/nn/_core_modules/_concatenator.py +++ b/src/phlower/nn/_core_modules/_concatenator.py @@ -3,6 +3,7 @@ import torch from typing_extensions import Self +from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils @@ -58,7 +59,7 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor] | None = None, + field_data: ISimulationField | None = None, **kwards, ) -> PhlowerTensor: """forward function which overloads torch.nn.Module diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index d0192c0..d197bca 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -2,6 +2,7 @@ import torch +from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _functions, _utils @@ -77,7 +78,7 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor], + field_data: ISimulationField, **kwards, ) -> PhlowerTensor: """forward function which overload torch.nn.Module @@ -92,7 +93,7 @@ def forward( PhlowerTensor: Tensor object """ - support = supports[self._support_name] + support = field_data[self._support_name] h = data.unique_item() for i in range(len(self._chains)): h = self._propagate(h, support) diff --git a/src/phlower/nn/_core_modules/_mlp.py b/src/phlower/nn/_core_modules/_mlp.py index 60f28e5..965cdfb 100644 --- a/src/phlower/nn/_core_modules/_mlp.py +++ b/src/phlower/nn/_core_modules/_mlp.py @@ -3,6 +3,7 @@ import torch from typing_extensions import Self +from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils @@ -72,6 +73,7 @@ def forward( self, data: IPhlowerTensorCollections, *, + field_data: ISimulationField | None = None, supports: dict[str, PhlowerTensor] | None = None, **kwards, ) -> PhlowerTensor: diff --git a/src/phlower/nn/_core_modules/_share.py b/src/phlower/nn/_core_modules/_share.py index 70fd7dd..f8fc7ad 100644 --- a/src/phlower/nn/_core_modules/_share.py +++ b/src/phlower/nn/_core_modules/_share.py @@ -2,6 +2,7 @@ import torch +from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._interface_module import ( @@ -56,7 +57,7 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor] | None = None, + field_data: ISimulationField | None = None, **kwards, ) -> PhlowerTensor: """forward function which overload torch.nn.Module @@ -77,4 +78,4 @@ def forward( "Please check that `resolve` function is called." ) - return self._reference.forward(data, supports=supports, **kwards) + return self._reference.forward(data, field_data=field_data, **kwards) diff --git a/src/phlower/nn/_group_module.py b/src/phlower/nn/_group_module.py index 4f74aad..741dd57 100644 --- a/src/phlower/nn/_group_module.py +++ b/src/phlower/nn/_group_module.py @@ -6,7 +6,7 @@ import torch from typing_extensions import Self -from phlower import PhlowerTensor +from phlower import ISimulationField, PhlowerTensor from phlower.collections.tensors import ( IPhlowerTensorCollections, phlower_tensor_collection, @@ -128,7 +128,7 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor] | None = None, + field_data: ISimulationField, **kwards, ) -> IPhlowerTensorCollections: results = phlower_tensor_collection({}) @@ -147,7 +147,7 @@ def forward( args = reduce_collections(node.get_received_args()) _module: PhlowerModuleAdapter = node.get_user_function() - _result = _module.forward(args, supports=supports, **kwards) + _result = _module.forward(args, field_data=field_data, **kwards) dag_modules.send(node.mut_name, _result) dag_modules.done(node.mut_name) diff --git a/src/phlower/nn/_interface_module.py b/src/phlower/nn/_interface_module.py index 412aaac..1fc2a48 100644 --- a/src/phlower/nn/_interface_module.py +++ b/src/phlower/nn/_interface_module.py @@ -5,6 +5,7 @@ from typing_extensions import Self +from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections from phlower.settings._module_settings import IPhlowerLayerParameters @@ -33,7 +34,7 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor] = None, + field_data: ISimulationField = None, ) -> PhlowerTensor: ... @abc.abstractmethod diff --git a/src/phlower/nn/_phlower_module_adpter.py b/src/phlower/nn/_phlower_module_adpter.py index 90f0d87..434e895 100644 --- a/src/phlower/nn/_phlower_module_adpter.py +++ b/src/phlower/nn/_phlower_module_adpter.py @@ -4,7 +4,7 @@ import torch -from phlower import PhlowerTensor +from phlower import ISimulationField from phlower.collections.tensors import ( IPhlowerTensorCollections, phlower_tensor_collection, @@ -92,16 +92,16 @@ def forward( self, data: IPhlowerTensorCollections, *, - supports: dict[str, PhlowerTensor], + field_data: ISimulationField, ) -> IPhlowerTensorCollections: inputs = phlower_tensor_collection( {key: data[key] for key in self._input_keys} ) if self._no_grad: with torch.no_grad(): - result = self._layer.forward(inputs, supports=supports) + result = self._layer.forward(inputs, field_data=field_data) else: - result = self._layer.forward(inputs, supports=supports) + result = self._layer.forward(inputs, field_data=field_data) return phlower_tensor_collection({self._output_key: result}) diff --git a/src/phlower/services/predictor/_predictor.py b/src/phlower/services/predictor/_predictor.py index c2827f9..062a675 100644 --- a/src/phlower/services/predictor/_predictor.py +++ b/src/phlower/services/predictor/_predictor.py @@ -55,9 +55,7 @@ def predict( for batch in data_loader: batch: LumpedTensorData - h = self._model.forward( - batch.x_data, supports=batch.sparse_supports - ) + h = self._model.forward(batch.x_data, supports=batch.field_data) yield h # HACK: Need to save h diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index 6e61f4e..d73cc70 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -125,7 +125,7 @@ def train( self._scheduled_optimizer.zero_grad() h = self._model.forward( - tr_batch.x_data, supports=tr_batch.sparse_supports + tr_batch.x_data, supports=tr_batch.field_data ) losses = loss_function.calculate( diff --git a/src/phlower/settings/__init__.py b/src/phlower/settings/__init__.py index cd3f3f1..bd10428 100644 --- a/src/phlower/settings/__init__.py +++ b/src/phlower/settings/__init__.py @@ -1,6 +1,6 @@ from phlower.settings._group_settings import GroupModuleSetting, ModuleSetting +from phlower.settings._model_setting import PhlowerModelSetting from phlower.settings._phlower_setting import ( - PhlowerModelSetting, PhlowerPredictorSetting, PhlowerSetting, ) diff --git a/src/phlower/settings/_model_setting.py b/src/phlower/settings/_model_setting.py new file mode 100644 index 0000000..127c6c3 --- /dev/null +++ b/src/phlower/settings/_model_setting.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import pydantic +from pipe import uniq +from typing_extensions import Self + +from phlower._base import PhysicalDimensionsClass +from phlower.settings._group_settings import GroupModuleSetting + + +class PhlowerFieldSetting(pydantic.BaseModel): + field_names: list[str] = pydantic.Field(default_factory=list, frozen=True) + """ + name of variables in simulation field which are treated as constant + in calculation. + For example, + + * support matrix for input graph structure + * boundary conditions + + """ + + @pydantic.model_validator(mode="after") + def check_duplicate_names(self) -> Self: + unique_names = list(self.field_names | uniq) + + if len(unique_names) != len(self.field_names): + raise ValueError( + "Duplicate name is found. A varaible name " + "in field setting must be unique." + ) + + return self + + +class PhlowerModelSetting(pydantic.BaseModel): + variable_dimensions: dict[str, PhysicalDimensionsClass] = pydantic.Field( + default_factory=lambda: {}, validate_default=True + ) + """ + dictionary which maps variable name to value + """ + + network: GroupModuleSetting + """ + define structure of neural network + """ + + fields: PhlowerFieldSetting = pydantic.Field( + default_factory=PhlowerFieldSetting + ) + """ + settings for fields dependent on your mesh or graph + """ + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict( + frozen=True, extra="forbid", arbitrary_types_allowed=True + ) + + def resolve(self) -> None: + """ + Resolve network relationship. Following items are checked. + + * network is a valid DAG, not cycled graph. + * input keywards in a module is defined + as output in the precedent module. + * set positive integer value to the value + which is defined as -1 in nodes. + """ + self.network.resolve(is_first=True) diff --git a/src/phlower/settings/_phlower_setting.py b/src/phlower/settings/_phlower_setting.py index d926ec2..188fd88 100644 --- a/src/phlower/settings/_phlower_setting.py +++ b/src/phlower/settings/_phlower_setting.py @@ -9,6 +9,7 @@ from phlower._base import PhysicalDimensionsClass from phlower.io import PhlowerYamlFile from phlower.settings._group_settings import GroupModuleSetting +from phlower.settings._model_setting import PhlowerModelSetting from phlower.settings._scaling_setting import PhlowerScalingSetting from phlower.settings._trainer_setting import PhlowerTrainerSetting from phlower.utils.enums import ModelSelectionType @@ -62,25 +63,6 @@ def read_yaml( return PhlowerSetting(**data) -class PhlowerModelSetting(pydantic.BaseModel): - variable_dimensions: dict[str, PhysicalDimensionsClass] = pydantic.Field( - default_factory=lambda: {}, validate_default=True - ) - """ - dictionary which maps variable name to value - """ - - network: GroupModuleSetting - """ - define structure of neural network - """ - - # special keyward to forbid extra fields in pydantic - model_config = pydantic.ConfigDict( - frozen=True, extra="forbid", arbitrary_types_allowed=True - ) - - @dc.dataclass(frozen=True, config=pydantic.ConfigDict(extra="forbid")) class PhlowerPredictorSetting: selection_mode: str diff --git a/tests/test_data/test_data_loader.py b/tests/test_data/test_data_loader.py index 08d1d18..dd2f4bf 100644 --- a/tests/test_data/test_data_loader.py +++ b/tests/test_data/test_data_loader.py @@ -70,7 +70,7 @@ def test__consider_batch_size( x_variable_names=["x0", "x1", "x2"], y_variable_names=["y0"], directories=directories, - support_names=["s0", "s1"], + field_names=["s0", "s1"], ) builder = DataLoaderBuilder( @@ -123,7 +123,7 @@ def test__consider_dimensions( x_variable_names=["x0", "x1", "x2"], y_variable_names=["y0"], directories=directories, - support_names=["s0"], + field_names=["s0"], ) builder = DataLoaderBuilder( @@ -149,10 +149,8 @@ def test__consider_dimensions( phydim = item.y_data[data_name].dimension.to_physics_dimension() assert phydim == desired[data_name] - for data_name in item.sparse_supports.keys(): - phydim = item.sparse_supports[ - data_name - ].dimension.to_physics_dimension() + for data_name in item.field_data.keys(): + phydim = item.field_data[data_name].dimension.to_physics_dimension() assert phydim == desired[data_name] @@ -181,7 +179,7 @@ def test__not_consider_dimensions( x_variable_names=["x0", "x1", "x2"], y_variable_names=["y0"], directories=directories, - support_names=["s0"], + field_names=["s0"], ) builder = DataLoaderBuilder( @@ -205,5 +203,5 @@ def test__not_consider_dimensions( for data_name in item.y_data.keys(): assert not item.y_data[data_name].has_dimension - for data_name in item.sparse_supports.keys(): - assert not item.sparse_supports[data_name].has_dimension + for data_name in item.field_data.keys(): + assert not item.field_data[data_name].has_dimension diff --git a/tests/test_data/test_datasets.py b/tests/test_data/test_datasets.py index 9798fdb..2105aca 100644 --- a/tests/test_data/test_datasets.py +++ b/tests/test_data/test_datasets.py @@ -17,19 +17,19 @@ def test__lazy_dataset_length( x_variable_names=["x0", "x1"], y_variable_names=["y0"], directories=directories, - support_names=["s0"], + field_names=["s0"], ) assert len(dataset) == desired @pytest.mark.parametrize( - "x_variable_names, y_variable_names, support_names, directory_names", + "x_variable_names, y_variable_names, field_names, directory_names", [(["x0", "x1", "x2"], ["y0"], ["s0", "s1"], ["data0", "data1", "data2"])], ) def test__lazy_dataset_getitem( x_variable_names, y_variable_names, - support_names, + field_names, directory_names, create_tmp_dataset, output_base_directory, @@ -39,7 +39,7 @@ def test__lazy_dataset_getitem( x_variable_names=x_variable_names, y_variable_names=y_variable_names, directories=directories, - support_names=support_names, + field_names=field_names ) assert len(dataset) > 1 @@ -60,15 +60,15 @@ def test__lazy_dataset_getitem( item.y_data[v_name].to_numpy(), desired[data_name][v_name] ) - for v_name in support_names: + for v_name in field_names: np.testing.assert_array_almost_equal( - item.sparse_supports[v_name].to_numpy().todense(), + item.field_data[v_name].to_numpy().todense(), desired[data_name][v_name].todense(), ) @pytest.mark.parametrize( - "x_variable_names, y_variable_names, support_names, directory_names", + "x_variable_names, y_variable_names, field_names, directory_names", [ (["x0", "x1", "x2"], None, ["s0", "s1"], ["data0", "data1", "data2"]), (["x0", "x1", "x2"], ["y3"], ["s0", "s1"], ["data0", "data1", "data2"]), @@ -77,7 +77,7 @@ def test__lazy_dataset_getitem( def test__lazy_dataset_getitem_when_no_ydata( x_variable_names, y_variable_names, - support_names, + field_names, directory_names, create_tmp_dataset, output_base_directory, @@ -87,7 +87,7 @@ def test__lazy_dataset_getitem_when_no_ydata( x_variable_names=x_variable_names, y_variable_names=y_variable_names, directories=directories, - support_names=support_names, + field_names=field_names, allow_no_y_data=True, ) assert len(dataset) > 1 diff --git a/tests/test_nn/test_core_modules/test_gcn.py b/tests/test_nn/test_core_modules/test_gcn.py index db767e4..6986d3a 100644 --- a/tests/test_nn/test_core_modules/test_gcn.py +++ b/tests/test_nn/test_core_modules/test_gcn.py @@ -38,7 +38,7 @@ def test__gcn(size, is_time_series): activations=["tanh", "identity"], ) - actual = model(phlower_tensors, supports=dict_supports) + actual = model(phlower_tensors, field_data=dict_supports) assert actual.shape == size assert actual.is_time_series == is_time_series diff --git a/tests/test_nn/test_group_module.py b/tests/test_nn/test_group_module.py index 5a39c92..7d05feb 100644 --- a/tests/test_nn/test_group_module.py +++ b/tests/test_nn/test_group_module.py @@ -4,6 +4,7 @@ import pytest import scipy.sparse as sp +from phlower._fields import SimulationField from phlower._base import phlower_array from phlower.collections import phlower_tensor_collection from phlower.nn import PhlowerGroupModule @@ -62,5 +63,6 @@ def test__forward_and_backward(yaml_file, input_n_feature, n_nodes): ) ) nodal_nadj = sparse_adj.to_phlower_tensor() + field_data = SimulationField(field_tensors={"support1": nodal_nadj}) - _ = group.forward(data=phlower_tensors, supports={"support1": nodal_nadj}) + _ = group.forward(data=phlower_tensors, field_data=field_data) diff --git a/tests/test_settings/data/groups/simple_module_with_simultion_field.yml b/tests/test_settings/data/groups/simple_module_with_simultion_field.yml new file mode 100644 index 0000000..310ebb3 --- /dev/null +++ b/tests/test_settings/data/groups/simple_module_with_simultion_field.yml @@ -0,0 +1,52 @@ +misc: + + tests: + MLP0: 10 + GCN0: 100 + +model: + variable_dimensions: + feature0: + + fields: + supports: + - name: support1 + boundaries: + - name: dirichlet_u + - name: neumann_p + + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_dim: 10 + - name: feature1 + n_dim: 12 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 100] + activations: ["Identity", "identity"] + + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 From 6cf09ee040a0edda8d83ed2c0ad2a2cdd5f039cf Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 14 Aug 2024 17:27:37 +0900 Subject: [PATCH 39/89] reflect mr --- src/phlower/nn/_core_modules/_functions.py | 72 ++---- src/phlower/nn/_core_modules/_mlp.py | 2 +- src/phlower/nn/_core_modules/_pinv_mlp.py | 39 +-- .../_similarity_equivariant_mlp.py | 25 +- src/phlower/nn/_core_modules/_utils.py | 23 ++ .../settings/_module_settings/_mlp_setting.py | 2 +- src/phlower/utils/enums.py | 2 +- .../test_similarity_equivariant_mlp.py | 222 ++++++++++++++++-- 8 files changed, 261 insertions(+), 126 deletions(-) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 3f9d09b..ebe558e 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -154,21 +154,10 @@ def contraction( raise PhlowerIncompatibleTensorError( "Cannot compute contraction between non-voxel and voxel.") - ret_is_time_series = False - if x.is_time_series: - time_x = "t" - ret_is_time_series = True - else: - time_x = "" - if y.is_time_series: - time_y = "t" - ret_is_time_series = True - else: - time_y = "" - if ret_is_time_series: - time_ret = "t" - else: - time_ret = "" + ret_is_time_series = x.is_time_series or y.is_time_series + time_x = "t" if x.is_time_series else "" + time_y = "t" if y.is_time_series else "" + time_ret = "t" if ret_is_time_series else "" # No need to consider y because they should be compatible if x.is_voxel: @@ -211,21 +200,10 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: x_rank = x.rank() y_rank = y.rank() - ret_is_time_series = False - if x.is_time_series: - x_time = "t" - ret_is_time_series = True - else: - x_time = "" - if y.is_time_series: - y_time = "t" - ret_is_time_series = True - else: - y_time = "" - if ret_is_time_series: - t_ret = "t" - else: - t_ret = "" + ret_is_time_series = x.is_time_series or y.is_time_series + time_x = "t" if x.is_time_series else "" + time_y = "t" if y.is_time_series else "" + time_ret = "t" if ret_is_time_series else "" # No need to consider y because they should be compatible if x.is_voxel: @@ -237,8 +215,8 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: x_vars = _availale_variables(x_rank) y_vars = _availale_variables(y_rank, start=x_rank) - equation = f"{x_time}{space}{x_vars}f,{y_time}{space}{y_vars}f->" \ - + f"{t_ret}{space}{x_vars}{y_vars}f" + equation = f"{time_x}{space}{x_vars}f,{time_y}{space}{y_vars}f->" \ + + f"{time_ret}{space}{x_vars}{y_vars}f" if x.dimension is None or y.dimension is None: dimension = None @@ -293,17 +271,11 @@ def apply_orthogonal_group( if rank == 0: return tensor + time = "t" if tensor.is_time_series else "" + space = "xyz" if tensor.is_voxel else "x" start_dim = 1 - if tensor.is_time_series: - time = "t" - start_dim += 1 - else: - time = "" - if tensor.is_voxel: - space = "xyz" - start_dim += 3 - else: - space = "x" + start_dim = start_dim + 1 if tensor.is_time_series else start_dim + start_dim = start_dim + 3 if tensor.is_voxel else start_dim s = _availale_variables(rank * 2) str_ortho = ','.join(a + b for a, b in zip(s[::2], s[1::2], strict=True)) @@ -321,18 +293,10 @@ def apply_orthogonal_group( def spatial_sum(tensor: IPhlowerTensor) -> IPhlowerTensor: """Compute sum over space.""" - if tensor.is_time_series: - time = "t" - start_space = 1 - else: - time = "" - start_space = 0 - if tensor.is_voxel: - space = "xyz" - space_width = 3 - else: - space = "x" - space_width = 1 + time = "t" if tensor.is_time_series else "" + space = "xyz" if tensor.is_voxel else "x" + start_space = 1 if tensor.is_time_series else 0 + space_width = 3 if tensor.is_voxel else 1 squeezed = einsum( f"{time}{space}...->{time}...", tensor, diff --git a/src/phlower/nn/_core_modules/_mlp.py b/src/phlower/nn/_core_modules/_mlp.py index 90397e9..d14f13b 100644 --- a/src/phlower/nn/_core_modules/_mlp.py +++ b/src/phlower/nn/_core_modules/_mlp.py @@ -46,7 +46,7 @@ def __init__( nodes: list[int], activations: list[str] | None = None, dropouts: list[float] | None = None, - bias: bool = False, + bias: bool = True, ) -> None: super().__init__() diff --git a/src/phlower/nn/_core_modules/_pinv_mlp.py b/src/phlower/nn/_core_modules/_pinv_mlp.py index d474912..69f9c41 100644 --- a/src/phlower/nn/_core_modules/_pinv_mlp.py +++ b/src/phlower/nn/_core_modules/_pinv_mlp.py @@ -11,10 +11,7 @@ IReadonlyReferenceGroup, ) from phlower.settings._module_settings import PInvMLPSetting -from phlower.utils.exceptions import ( - NotFoundReferenceModuleError, - PhlowerInvalidActivationError, -) +from phlower.utils.exceptions import NotFoundReferenceModuleError class PInvMLP(IPhlowerCoreModule, torch.nn.Module): @@ -60,27 +57,11 @@ def resolve( def _initialize(self): """Initialize parameters and activations after resolve() is called.""" - self._activation_names = [ - self._inverse_activation_name(a) - for a in self._reference._activations[::-1]] self._activations = [ - _utils.ActivationSelector.select(name) - for name in self._activation_names] + _utils.ActivationSelector.select_inverse(name) + for name in self._reference._activations[::-1]] self._chains = self._init_pinv_chains() - def _inverse_activation_name(self, activation_name): - if activation_name == "identity": - return "identity" - if activation_name == "leaky_relu0p5": - return "inversed_leaky_relu0p5" - if activation_name == "smooth_leaky_relu": - return "inversed_smooth_leaky_relu" - if activation_name == "tanh": - return "truncated_atanh" - - raise PhlowerInvalidActivationError( - f"Cannot pinv for {activation_name}") - def _init_pinv_chains(self): name = self._reference.__class__.__name__ if name in ["MLP", "Proportional"]: @@ -89,8 +70,8 @@ def _init_pinv_chains(self): raise ValueError(f"Unsupported reference class: {name}") def _init_pinv_mlp_chains( - self, chains: _utils.ExtendedLinearList, option=None): - return [PInvLinear(c, option=option) for c in chains._linears[::-1]] + self, chains: _utils.ExtendedLinearList): + return [PInvLinear(c) for c in chains._linears[::-1]] def forward( self, @@ -128,10 +109,9 @@ def forward( class PInvLinear(torch.nn.Module): - def __init__(self, ref_linear: torch.nn.Linear, option: str | None = None): + def __init__(self, ref_linear: torch.nn.Linear): super().__init__() self.ref_linear = ref_linear - self.option = option return def forward(self, x): @@ -141,10 +121,7 @@ def forward(self, x): @property def weight(self): """Return pseudo inversed weight.""" - if self.option is None: - w = self.ref_linear.weight - else: - raise ValueError(f"Unexpected option: {self.option}") + w = self.ref_linear.weight return torch.pinverse(w) @property @@ -153,4 +130,4 @@ def bias(self): if self.ref_linear.bias is None: return 0 else: - return - self.ref.bias + return - self.ref_linear.bias diff --git a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py index 9c035f3..0de0035 100644 --- a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py @@ -116,39 +116,36 @@ def forward( data: IPhlowerTensorCollections, *, supports: dict[str, PhlowerTensor] | None = None, + dict_scales: dict[str, PhlowerTensor | float] | None = None, **kwards, ) -> PhlowerTensor: """forward function which overloads torch.nn.Module Args: data (IPhlowerTensorCollections): - Data with the keys in tensor or PhysicalDimensionSymbolType, - with tensor is the input tensor to be computed. - PhysicalDimensionSymbolType keys denote physical scales. - No need to input all dimensions, but at least one scale - information should be input. + Data which receives from predecessors. supports (dict[str, PhlowerTensor], optional): Graph object. Defaults to None. Returns: PhlowerTensor: Tensor object """ - # TODO: Check if we can assume key convention - h = data.pop("tensor") + h = data.unique_item() if not h.has_dimension: raise PhlowerDimensionRequiredError("Dimension is required") - - if len(data) < 1: - raise PhlowerInvalidArgumentsError("Scale inputs are required") + if dict_scales is None or len(data) < 1: + raise PhlowerInvalidArgumentsError( + f"Scale inputs are required. Given: {dict_scales}, {kwards}") if not np.all([ - PhysicalDimensionSymbolType.is_exist(k) for k in data.keys()]): + PhysicalDimensionSymbolType.is_exist(k) + for k in dict_scales.keys()]): raise PhlowerInvalidArgumentsError( - "keys should be in PhysicalDimensionSymbolType. " - f"Given: {data.keys()}") + "keys in dict_scales should be in " + "PhysicalDimensionSymbolType. Given: {dict_scales.keys()}") dict_dimension = h.dimension.to_dict() dict_scales = { - k: v**dict_dimension[k] for k, v in data.items()} + k: v**dict_dimension[k] for k, v in dict_scales.items()} # Make h dimensionless for v in dict_scales.values(): diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index e625074..96afb8e 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -4,6 +4,7 @@ from phlower._base.tensors import PhlowerTensor from phlower.nn._core_modules import _functions +from phlower.utils.exceptions import PhlowerInvalidActivationError class ExtendedLinearList(torch.nn.Module): @@ -96,6 +97,28 @@ def select(name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: name = "identity" return ActivationSelector._REGISTERED_ACTIVATIONS[name] + @staticmethod + def select_inverse( + name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: + if name is None: + name = "identity" + return ActivationSelector._REGISTERED_ACTIVATIONS[ + ActivationSelector._inverse_activation_name(name)] + + @staticmethod + def _inverse_activation_name(activation_name: str) -> str: + if activation_name == "identity": + return "identity" + if activation_name == "leaky_relu0p5": + return "inversed_leaky_relu0p5" + if activation_name == "smooth_leaky_relu": + return "inversed_smooth_leaky_relu" + if activation_name == "tanh": + return "truncated_atanh" + + raise PhlowerInvalidActivationError( + f"Cannot inverse for {activation_name}") + @staticmethod def is_exists(name: str) -> bool: return name in ActivationSelector._REGISTERED_ACTIVATIONS diff --git a/src/phlower/settings/_module_settings/_mlp_setting.py b/src/phlower/settings/_module_settings/_mlp_setting.py index b4d02ce..d72fc50 100644 --- a/src/phlower/settings/_module_settings/_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_mlp_setting.py @@ -16,7 +16,7 @@ class MLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): ) # This property only overwritten when resolving. activations: list[str] = Field(default_factory=lambda: [], frozen=True) dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) - bias: bool = Field(False, frozen=True) + bias: bool = Field(True, frozen=True) def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: diff --git a/src/phlower/utils/enums.py b/src/phlower/utils/enums.py index 48b3e39..ade9d15 100644 --- a/src/phlower/utils/enums.py +++ b/src/phlower/utils/enums.py @@ -52,7 +52,7 @@ class PhysicalDimensionSymbolType(Enum): T = 0 # time L = 1 # length M = 2 # mass - I = 3 # electric current # NOQA + I = 3 # electric current # noqa: E741 Theta = 4 # thermodynamic temperature N = 5 # amount of substance J = 6 # luminous intensity diff --git a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py index 3508e7c..603144f 100644 --- a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py +++ b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py @@ -56,21 +56,198 @@ def test__can_call_parameters(): @pytest.mark.parametrize( "norm_function_name", ["identity", "sqrt"]) @pytest.mark.parametrize( - "disable_en_equivariance", [False, True]) + "centering", [False, True]) +def test__similarity_equivariance( + size, is_time_series, is_voxel, activation, n_output_feature, + dimension, norm_function_name, centering, +): + orthogonal_tensor = PhlowerTensor( + torch.tensor(ortho_group.rvs(3).astype(np.float32))) + dict_scaling_factor = { + k: + np.random.rand() * 10 for k, v in dimension.items()} + scaling_factor = np.prod( + [dict_scaling_factor[k]**v for k, v in dimension.items()]) + + create_linear_weight = size[-1] != n_output_feature + model = SimilarityEquivariantMLP( + nodes=[size[-1], n_output_feature, n_output_feature], + activations=[activation, activation], + create_linear_weight=create_linear_weight, + norm_function_name=norm_function_name, + disable_en_equivariance=False, + centering=centering, + ) + + t = phlower_tensor( + torch.rand(*size) * 2 - 1., dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + dict_tensor = {"tensor": t} + dict_scales = { + k: phlower_tensor( + torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) + for k, v in dimension.items() + } + + ts = phlower_tensor_collection(dict_tensor) + actual_tensor = _functions.apply_orthogonal_group( + orthogonal_tensor, + model(ts, dict_scales=dict_scales)) * scaling_factor + actual = actual_tensor.to_numpy() + + dict_transformed = {'tensor': _functions.apply_orthogonal_group( + orthogonal_tensor, t) * scaling_factor} + dict_scaled_scales = { + k: v * dict_scaling_factor[k] for k, v in dict_scales.items()} + + transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) + desired = model( + transformed_phlower_tensors, dict_scales=dict_scaled_scales).to_numpy() + + scale = np.max(np.abs(desired)) + # Test equivariance + np.testing.assert_almost_equal(actual / scale, desired / scale, decimal=2) + + # Test dimension is kept + for k, v in actual_tensor.dimension.to_dict().items(): + np.testing.assert_almost_equal(v, dimension[k]) + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel", + [ + ((10, 1), False, False), + ((10, 16), False, False), + ((10, 3, 16), False, False), + ((4, 10, 1), True, False), + ((4, 10, 16), True, False), + ((4, 10, 3, 16), True, False), + ((10, 10, 10, 1), False, True), + ((10, 10, 10, 16), False, True), + ((10, 10, 10, 3, 16), False, True), + ((4, 10, 10, 10, 1), True, True), + ((4, 10, 10, 10, 16), True, True), + ((4, 10, 10, 10, 3, 16), True, True), + ], +) +@pytest.mark.parametrize("activation", ["identity", "tanh", "leaky_relu0p5"]) +@pytest.mark.parametrize("n_output_feature", [1, 16, 32]) +@pytest.mark.parametrize( + "dimension", + [ + # Dimensionless + {"T": 0, "L": 0, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Velocity + {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Mass density + {"T": 0, "L": - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Momentum density + {"T": -1, "L": 1 - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + ], +) +@pytest.mark.parametrize( + "norm_function_name", ["identity", "sqrt"]) +@pytest.mark.parametrize( + "centering", [False, True]) +def test__similarity_invariance( + size, is_time_series, is_voxel, activation, n_output_feature, + dimension, norm_function_name, centering, +): + orthogonal_tensor = PhlowerTensor( + torch.tensor(ortho_group.rvs(3).astype(np.float32))) + dict_scaling_factor = { + k: + np.random.rand() * 10 for k, v in dimension.items()} + scaling_factor = np.prod( + [dict_scaling_factor[k]**v for k, v in dimension.items()]) + + create_linear_weight = size[-1] != n_output_feature + model = SimilarityEquivariantMLP( + nodes=[size[-1], n_output_feature, n_output_feature], + activations=[activation, activation], + create_linear_weight=create_linear_weight, + norm_function_name=norm_function_name, + disable_en_equivariance=False, + centering=centering, invariant=True, + ) + + t = phlower_tensor( + torch.rand(*size) * 2 - 1., dimension=dimension, + is_time_series=is_time_series, is_voxel=is_voxel) + dict_tensor = {"tensor": t} + dict_scales = { + k: phlower_tensor( + torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) + for k, v in dimension.items() + } + + ts = phlower_tensor_collection(dict_tensor) + actual_tensor = _functions.apply_orthogonal_group( + orthogonal_tensor, model(ts, dict_scales=dict_scales)) + actual = actual_tensor.to_numpy() + + dict_transformed = {'tensor': _functions.apply_orthogonal_group( + orthogonal_tensor, t) * scaling_factor} + dict_scaled_scales = { + k: v * dict_scaling_factor[k] for k, v in dict_scales.items()} + + transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) + desired = model( + transformed_phlower_tensors, dict_scales=dict_scaled_scales).to_numpy() + + scale = np.max(np.abs(desired)) + # Test equivariance + np.testing.assert_almost_equal(actual / scale, desired / scale, decimal=2) + + # Test dimensionless in case of invariant + for v in actual_tensor.dimension.to_dict().values(): + np.testing.assert_almost_equal(v, 0.) + + +@pytest.mark.parametrize( + "size, is_time_series, is_voxel", + [ + ((10, 1), False, False), + ((10, 16), False, False), + ((10, 3, 16), False, False), + ((4, 10, 1), True, False), + ((4, 10, 16), True, False), + ((4, 10, 3, 16), True, False), + ((10, 10, 10, 1), False, True), + ((10, 10, 10, 16), False, True), + ((10, 10, 10, 3, 16), False, True), + ((4, 10, 10, 10, 1), True, True), + ((4, 10, 10, 10, 16), True, True), + ((4, 10, 10, 10, 3, 16), True, True), + ], +) +@pytest.mark.parametrize("activation", ["identity", "tanh", "leaky_relu0p5"]) +@pytest.mark.parametrize("n_output_feature", [1, 16, 32]) +@pytest.mark.parametrize( + "dimension", + [ + # Dimensionless + {"T": 0, "L": 0, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Velocity + {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Mass density + {"T": 0, "L": - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + # Momentum density + {"T": -1, "L": 1 - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + ], +) +@pytest.mark.parametrize( + "norm_function_name", ["identity", "sqrt"]) @pytest.mark.parametrize( "centering", [False, True]) @pytest.mark.parametrize( "invariant", [False, True]) -def test__similarity_equivariance( +def test__scaling_equivariance( size, is_time_series, is_voxel, activation, n_output_feature, - dimension, norm_function_name, disable_en_equivariance, + dimension, norm_function_name, centering, invariant, ): - if disable_en_equivariance: - orthogonal_tensor = PhlowerTensor(torch.eye(3)) - else: - orthogonal_tensor = PhlowerTensor( - torch.tensor(ortho_group.rvs(3).astype(np.float32))) + orthogonal_tensor = PhlowerTensor(torch.eye(3)) dict_scaling_factor = { k: np.random.rand() * 10 for k, v in dimension.items()} @@ -83,44 +260,40 @@ def test__similarity_equivariance( activations=[activation, activation], create_linear_weight=create_linear_weight, norm_function_name=norm_function_name, - disable_en_equivariance=disable_en_equivariance, + disable_en_equivariance=True, centering=centering, invariant=invariant, ) t = phlower_tensor( - torch.randn(*size), dimension=dimension, + torch.rand(*size) * 2 - 1., dimension=dimension, is_time_series=is_time_series, is_voxel=is_voxel) dict_tensor = {"tensor": t} - scales = { + dict_scales = { k: phlower_tensor( torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) for k, v in dimension.items() - if v != 0 } - if len(scales) == 0: - scales = {"L": phlower_tensor( - torch.rand(1) * 10, dimension=PhysicalDimensions({"L": 1}))} - dict_tensor.update(scales) ts = phlower_tensor_collection(dict_tensor) if invariant: actual_tensor = _functions.apply_orthogonal_group( - orthogonal_tensor, model(ts)) * 1. + orthogonal_tensor, model(ts, dict_scales=dict_scales)) * 1. else: actual_tensor = _functions.apply_orthogonal_group( - orthogonal_tensor, model(ts)) * scaling_factor + orthogonal_tensor, + model(ts, dict_scales=dict_scales)) * scaling_factor actual = actual_tensor.to_numpy() dict_transformed = {'tensor': _functions.apply_orthogonal_group( orthogonal_tensor, t) * scaling_factor} dict_scaled_scales = { - k: v * dict_scaling_factor[k] for k, v in scales.items()} - dict_transformed.update(dict_scaled_scales) + k: v * dict_scaling_factor[k] for k, v in dict_scales.items()} transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) - desired = model(transformed_phlower_tensors).to_numpy() + desired = model( + transformed_phlower_tensors, dict_scales=dict_scaled_scales).to_numpy() - scale = np.max(desired) + scale = np.max(np.abs(desired)) # Test equivariance np.testing.assert_almost_equal(actual / scale, desired / scale, decimal=2) @@ -140,9 +313,10 @@ def test__similarity_equivariance_no_dimension(): t = phlower_tensor(torch.rand(10, 3, 8), dimension=None) t_scale = phlower_tensor( torch.tensor(0.1), dimension=PhysicalDimensions({"T": 1})) - ts = phlower_tensor_collection({"tensor": t, "T": t_scale}) + dict_scales = {"T": t_scale} + ts = phlower_tensor_collection({"tensor": t}) with pytest.raises(PhlowerDimensionRequiredError): - model(ts) + model(ts, dict_scales=dict_scales) def test__similarity_equivariance_no_scale_input(): model = SimilarityEquivariantMLP(nodes=[8, 8]) From 8ee2727a0333da242404574fc452b6d784dda80f Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 14 Aug 2024 17:28:37 +0900 Subject: [PATCH 40/89] fix lint warnings --- src/phlower/_fields/_simulation_field.py | 2 +- src/phlower/data/_datasets.py | 2 +- src/phlower/nn/_group_module.py | 2 +- src/phlower/settings/_phlower_setting.py | 2 -- tests/test_data/test_datasets.py | 2 +- tests/test_nn/test_group_module.py | 2 +- 6 files changed, 5 insertions(+), 7 deletions(-) diff --git a/src/phlower/_fields/_simulation_field.py b/src/phlower/_fields/_simulation_field.py index a572a3c..7da4c05 100644 --- a/src/phlower/_fields/_simulation_field.py +++ b/src/phlower/_fields/_simulation_field.py @@ -22,7 +22,7 @@ def __init__( self._field_tensors = field_tensors if batch_info is None: - batch_info = {} + batch_info = {} self._batch_info = batch_info def keys(self): diff --git a/src/phlower/data/_datasets.py b/src/phlower/data/_datasets.py index 8a0cdad..b62f6e3 100644 --- a/src/phlower/data/_datasets.py +++ b/src/phlower/data/_datasets.py @@ -25,7 +25,7 @@ def __init__( *, field_names: list[str] = None, allow_no_y_data: bool = False, - decrypt_key: bytes | None = None + decrypt_key: bytes | None = None, ): self._x_variable_names = x_variable_names if y_variable_names is None: diff --git a/src/phlower/nn/_group_module.py b/src/phlower/nn/_group_module.py index 741dd57..e8dd421 100644 --- a/src/phlower/nn/_group_module.py +++ b/src/phlower/nn/_group_module.py @@ -6,7 +6,7 @@ import torch from typing_extensions import Self -from phlower import ISimulationField, PhlowerTensor +from phlower import ISimulationField from phlower.collections.tensors import ( IPhlowerTensorCollections, phlower_tensor_collection, diff --git a/src/phlower/settings/_phlower_setting.py b/src/phlower/settings/_phlower_setting.py index 188fd88..8209563 100644 --- a/src/phlower/settings/_phlower_setting.py +++ b/src/phlower/settings/_phlower_setting.py @@ -6,9 +6,7 @@ from pydantic import dataclasses as dc from typing_extensions import Self -from phlower._base import PhysicalDimensionsClass from phlower.io import PhlowerYamlFile -from phlower.settings._group_settings import GroupModuleSetting from phlower.settings._model_setting import PhlowerModelSetting from phlower.settings._scaling_setting import PhlowerScalingSetting from phlower.settings._trainer_setting import PhlowerTrainerSetting diff --git a/tests/test_data/test_datasets.py b/tests/test_data/test_datasets.py index 2105aca..eed0963 100644 --- a/tests/test_data/test_datasets.py +++ b/tests/test_data/test_datasets.py @@ -39,7 +39,7 @@ def test__lazy_dataset_getitem( x_variable_names=x_variable_names, y_variable_names=y_variable_names, directories=directories, - field_names=field_names + field_names=field_names, ) assert len(dataset) > 1 diff --git a/tests/test_nn/test_group_module.py b/tests/test_nn/test_group_module.py index 7d05feb..dd90f5a 100644 --- a/tests/test_nn/test_group_module.py +++ b/tests/test_nn/test_group_module.py @@ -4,8 +4,8 @@ import pytest import scipy.sparse as sp -from phlower._fields import SimulationField from phlower._base import phlower_array +from phlower._fields import SimulationField from phlower.collections import phlower_tensor_collection from phlower.nn import PhlowerGroupModule from phlower.settings import PhlowerSetting From f48bc75c725e6de25c0754378f13ec040d089c18 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 14 Aug 2024 17:36:33 +0900 Subject: [PATCH 41/89] fix argument of trainer class --- src/phlower/services/predictor/_predictor.py | 4 ++-- src/phlower/services/trainer/_trainer.py | 9 +++++---- tests/e2e_tests/data/train.yml | 6 ++++-- tests/e2e_tests/data/train_batch_size.yml | 6 ++++-- 4 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/phlower/services/predictor/_predictor.py b/src/phlower/services/predictor/_predictor.py index 062a675..0562c9b 100644 --- a/src/phlower/services/predictor/_predictor.py +++ b/src/phlower/services/predictor/_predictor.py @@ -40,7 +40,7 @@ def predict( dataset = LazyPhlowerDataset( x_variable_names=self._model_setting.network.get_input_keys(), y_variable_names=self._model_setting.network.get_output_keys(), - support_names=self._model_setting.network.support_names, + field_names=self._model_setting.fields.field_names, directories=preprocessed_directories, ) @@ -55,7 +55,7 @@ def predict( for batch in data_loader: batch: LumpedTensorData - h = self._model.forward(batch.x_data, supports=batch.field_data) + h = self._model.forward(batch.x_data, field_data=batch.field_data) yield h # HACK: Need to save h diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index d73cc70..39627ad 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -84,13 +84,13 @@ def train( train_dataset = LazyPhlowerDataset( x_variable_names=self._model_setting.network.get_input_keys(), y_variable_names=self._model_setting.network.get_output_keys(), - support_names=self._model_setting.network.support_names, + field_names=self._model_setting.fields.field_names, directories=train_directories, ) validation_dataset = LazyPhlowerDataset( x_variable_names=self._model_setting.network.get_input_keys(), y_variable_names=self._model_setting.network.get_output_keys(), - support_names=self._model_setting.network.support_names, + field_names=self._model_setting.fields.field_names, directories=validation_directories, ) @@ -125,7 +125,7 @@ def train( self._scheduled_optimizer.zero_grad() h = self._model.forward( - tr_batch.x_data, supports=tr_batch.field_data + tr_batch.x_data, field_data=tr_batch.field_data ) losses = loss_function.calculate( @@ -145,8 +145,9 @@ def train( self._model.eval() for val_batch in validation_loader: with torch.no_grad(): + val_batch: LumpedTensorData h = self._model.forward( - val_batch.x_data, supports=val_batch.sparse_supports + val_batch.x_data, field_data=val_batch.field_data ) val_losses = loss_function.calculate( h, diff --git a/tests/e2e_tests/data/train.yml b/tests/e2e_tests/data/train.yml index 9b41c62..41321d0 100644 --- a/tests/e2e_tests/data/train.yml +++ b/tests/e2e_tests/data/train.yml @@ -17,6 +17,10 @@ model: nodal_last_u: {"L": 1, "T": -1} nodal_nadj: {} + fields: + field_names: + - "nodal_nadj" + network: nn_type: GROUP name: DEMO @@ -28,8 +32,6 @@ model: - name: nodal_last_u n_dim: 1 - support_names: ["nodal_nadj"] - modules: - nn_type: MLP name: MLP0 diff --git a/tests/e2e_tests/data/train_batch_size.yml b/tests/e2e_tests/data/train_batch_size.yml index 409c59e..c720955 100644 --- a/tests/e2e_tests/data/train_batch_size.yml +++ b/tests/e2e_tests/data/train_batch_size.yml @@ -16,6 +16,10 @@ model: nodal_last_u: {"L": 1, "T": -1} nodal_nadj: {} + fields: + field_names: + - "nodal_nadj" + network: nn_type: GROUP name: DEMO @@ -27,8 +31,6 @@ model: - name: nodal_last_u n_dim: 1 - support_names: ["nodal_nadj"] - modules: - nn_type: MLP name: MLP0 From ae419dc438666cdf37a86a0240b9392bd77d34f7 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 14 Aug 2024 17:41:33 +0900 Subject: [PATCH 42/89] debug conflict resolution --- src/phlower/nn/_core_modules/_functions.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index a336afe..37fc9ca 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -72,7 +72,6 @@ def spmm( for _ in range(repeat): h = torch.sparse.mm(sparse, h) return h.rearrange( -<<<<<<< HEAD pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, **dict_shape) From a3a922389dabe270cadf25ba79d75bf00a33ddb6 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 14 Aug 2024 17:43:44 +0900 Subject: [PATCH 43/89] fix sphinx tutorial settings yaml --- tutorials/basic_usages/sample_data/e2e/setting.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tutorials/basic_usages/sample_data/e2e/setting.yml b/tutorials/basic_usages/sample_data/e2e/setting.yml index fa95bdf..c0acf1a 100644 --- a/tutorials/basic_usages/sample_data/e2e/setting.yml +++ b/tutorials/basic_usages/sample_data/e2e/setting.yml @@ -26,7 +26,10 @@ scaling: training: batch_size: 1 random_seed: 0 - lr: 0.0001 + optimizer_setting: + optimizer: SGD + parameters: + lr: 0.0001 loss_setting: name2loss: nodal_last_u: "mse" From dd799fed75b249430ca197abd99dadb274308cff Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 14 Aug 2024 18:03:26 +0900 Subject: [PATCH 44/89] update for lint --- .../_base/tensors/_dimension_tensor.py | 12 +- src/phlower/_base/tensors/_phlower_tensor.py | 3 +- src/phlower/nn/__init__.py | 5 +- src/phlower/nn/_core_modules/__init__.py | 5 +- .../nn/_core_modules/_en_equivariant_mlp.py | 9 +- src/phlower/nn/_core_modules/_functions.py | 107 ++- src/phlower/nn/_core_modules/_pinv_mlp.py | 12 +- src/phlower/nn/_core_modules/_proportional.py | 4 +- .../_similarity_equivariant_mlp.py | 39 +- src/phlower/nn/_core_modules/_utils.py | 9 +- .../settings/_module_settings/__init__.py | 5 +- .../_en_equivariant_mlp_setting.py | 3 +- .../_similarity_equivariant_mlp_setting.py | 9 +- src/phlower/utils/exceptions.py | 4 + .../test_tensors/test__phlower_tensor.py | 8 +- .../test_en_equivariant_mlp.py | 25 +- .../test_core_modules/test_functions.py | 766 +++++++++++------- .../test_core_modules/test_identity.py | 6 +- tests/test_nn/test_core_modules/test_pinv.py | 4 +- .../test_core_modules/test_proportional.py | 14 +- .../test_similarity_equivariant_mlp.py | 195 +++-- 21 files changed, 783 insertions(+), 461 deletions(-) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 08488ab..e8bdfd0 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -150,7 +150,8 @@ def add(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): if inputs != other: raise DimensionIncompatibleError( "Add operation for different physical dimensions is not " - "allowed.") + "allowed." + ) return PhlowerDimensionTensor(inputs._tensor) @@ -163,7 +164,8 @@ def sub(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): if inputs != other: raise DimensionIncompatibleError( "Sub operation for different physical dimensions is not " - "allowed.") + "allowed." + ) return PhlowerDimensionTensor(inputs._tensor) @@ -294,7 +296,8 @@ def concatenate(inputs, *args, **kwards): def tanh(tensor: PhlowerDimensionTensor): if not tensor.is_dimensionless: raise DimensionIncompatibleError( - f"Should be dimensionless to apply tanh but {tensor}") + f"Should be dimensionless to apply tanh but {tensor}" + ) return tensor @@ -302,5 +305,6 @@ def tanh(tensor: PhlowerDimensionTensor): def leaky_relu(tensor: PhlowerDimensionTensor, *args, **kwargs): if not tensor.is_dimensionless: raise DimensionIncompatibleError( - f"Should be dimensionless to apply leaky_relu but {tensor}") + f"Should be dimensionless to apply leaky_relu but {tensor}" + ) return tensor diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index b13c1a9..ac68b33 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -106,7 +106,8 @@ def __init__( ): if not isinstance(tensor, torch.Tensor): raise PhlowerTypeError( - f"Expect torch.Tensor but {tensor.__class__} was fed") + f"Expect torch.Tensor but {tensor.__class__} was fed" + ) self._tensor = tensor self._dimension_tensor = dimension_tensor self._is_time_series = is_time_series diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 5f4d31f..0259afb 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -6,7 +6,8 @@ from phlower.nn._core_modules._pinv_mlp import PInvMLP from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share -from phlower.nn._core_modules \ - ._similarity_equivariant_mlp import SimilarityEquivariantMLP +from phlower.nn._core_modules._similarity_equivariant_mlp import ( + SimilarityEquivariantMLP, +) from phlower.nn._group_module import PhlowerGroupModule from phlower.nn._interface_module import IPhlowerCoreModule diff --git a/src/phlower/nn/_core_modules/__init__.py b/src/phlower/nn/_core_modules/__init__.py index 5774589..78f8ed6 100644 --- a/src/phlower/nn/_core_modules/__init__.py +++ b/src/phlower/nn/_core_modules/__init__.py @@ -10,8 +10,9 @@ if True: # NOTE: Import advanced models after from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP - from phlower.nn._core_modules \ - ._similarity_equivariant_mlp import SimilarityEquivariantMLP + from phlower.nn._core_modules._similarity_equivariant_mlp import ( + SimilarityEquivariantMLP, + ) _all_models: list[IPhlowerCoreModule] = [ Concatenator, diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py index 96aa9e8..5d67649 100644 --- a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py @@ -58,7 +58,8 @@ def __init__( dropouts = [] self._chains = _utils.ExtendedLinearList( - nodes=nodes, activations=activations, dropouts=dropouts, bias=bias) + nodes=nodes, activations=activations, dropouts=dropouts, bias=bias + ) self._nodes = nodes self._activations = activations self._create_linear_weight = create_linear_weight @@ -66,14 +67,16 @@ def __init__( self._linear_weight = self._init_linear_weight() self._norm_function = _utils.ActivationSelector.select( - self._norm_function_name) + self._norm_function_name + ) def _init_linear_weight(self): if not self._create_linear_weight: if self._nodes[0] != self._nodes[-1]: raise ValueError( "First and last nodes are different. " - "Set create_linear_weight True.") + "Set create_linear_weight True." + ) return Identity() return Proportional([self._nodes[0], self._nodes[-1]]) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 37fc9ca..09eab0e 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -23,8 +23,8 @@ def inversed_leaky_relu0p5(x): def truncated_atanh(x, epsilon=1e-8): """Inverse tanh with truncating values >=1 or <=-1.""" - x[x>=1. - epsilon] = 1. - epsilon - x[x<=-1. + epsilon] = -1. + epsilon + x[x >= 1.0 - epsilon] = 1.0 - epsilon + x[x <= -1.0 + epsilon] = -1.0 + epsilon return torch.atanh(x) @@ -43,7 +43,8 @@ def inverse(self, x): return ( self.a * x - torch.sqrt( - (self.a - 1)**2 * (2 * self.a * self.b - self.b + x**2)) + (self.a - 1) ** 2 * (2 * self.a * self.b - self.b + x**2) + ) ) / (2 * self.a - 1) def derivative(self, x): @@ -72,22 +73,27 @@ def spmm( for _ in range(repeat): h = torch.sparse.mm(sparse, h) return h.rearrange( - pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, - **dict_shape) + pattern, + is_time_series=x.is_time_series, + is_voxel=x.is_voxel, + **dict_shape, + ) def einsum( - equation, *args: list[IPhlowerTensor], - dimension: ( - PhysicalDimensions - | PhlowerDimensionTensor - | torch.Tensor - | dict[str, float] - | list[float] - | tuple[float] - | None - ) = None, - is_time_series: bool = False, is_voxel: bool = False, + equation, + *args: list[IPhlowerTensor], + dimension: ( + PhysicalDimensions + | PhlowerDimensionTensor + | torch.Tensor + | dict[str, float] + | list[float] + | tuple[float] + | None + ) = None, + is_time_series: bool = False, + is_voxel: bool = False, ) -> IPhlowerTensor: """ Compute einsum for phlower tensors. @@ -111,15 +117,17 @@ def einsum( Resultant tensor """ try: - ret_tensor = torch.einsum( - equation, [a.to_tensor() for a in args]) + ret_tensor = torch.einsum(equation, [a.to_tensor() for a in args]) except RuntimeError as e: raise PhlowerIncompatibleTensorError( - f"{e}\n" - f"{equation}, {[a.shape for a in args]}") from e + f"{e}\n" f"{equation}, {[a.shape for a in args]}" + ) from e return phlower_tensor( - ret_tensor, dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + ret_tensor, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) def _availale_variables(length: int, start: int = 0) -> str: @@ -128,11 +136,12 @@ def _availale_variables(length: int, start: int = 0) -> str: if length > len(available_variables): raise ValueError(f"Required length too long: {length}") - return available_variables[start:start+length] + return available_variables[start : start + length] def contraction( - x: IPhlowerTensor, y: IPhlowerTensor | None = None) -> IPhlowerTensor: + x: IPhlowerTensor, y: IPhlowerTensor | None = None +) -> IPhlowerTensor: """ Compute the tensor contraction. @@ -151,7 +160,8 @@ def contraction( if x.is_voxel != y.is_voxel: raise PhlowerIncompatibleTensorError( - "Cannot compute contraction between non-voxel and voxel.") + "Cannot compute contraction between non-voxel and voxel." + ) ret_is_time_series = x.is_time_series or y.is_time_series time_x = "t" if x.is_time_series else "" @@ -176,8 +186,12 @@ def contraction( return einsum( f"{time_x}{space}...{unresolved}f,{time_y}{space}...f->" f"{time_ret}{space}{unresolved}f", - x, y, dimension=dimension, - is_time_series=ret_is_time_series, is_voxel=is_voxel) + x, + y, + dimension=dimension, + is_time_series=ret_is_time_series, + is_voxel=is_voxel, + ) def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: @@ -194,7 +208,8 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: """ if x.is_voxel != y.is_voxel: raise PhlowerIncompatibleTensorError( - "Cannot compute contraction between non-voxel and voxel.") + "Cannot compute contraction between non-voxel and voxel." + ) x_rank = x.rank() y_rank = y.rank() @@ -214,8 +229,10 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: x_vars = _availale_variables(x_rank) y_vars = _availale_variables(y_rank, start=x_rank) - equation = f"{time_x}{space}{x_vars}f,{time_y}{space}{y_vars}f->" \ + equation = ( + f"{time_x}{space}{x_vars}f,{time_y}{space}{y_vars}f->" + f"{time_ret}{space}{x_vars}{y_vars}f" + ) if x.dimension is None or y.dimension is None: dimension = None @@ -223,12 +240,18 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: dimension = x.dimension * y.dimension return einsum( - equation, x, y, dimension=dimension, - is_time_series=ret_is_time_series, is_voxel=is_voxel) + equation, + x, + y, + dimension=dimension, + is_time_series=ret_is_time_series, + is_voxel=is_voxel, + ) def tensor_times_scalar( - tensor: IPhlowerTensor, scalar: int | float | IPhlowerTensor): + tensor: IPhlowerTensor, scalar: int | float | IPhlowerTensor +): """ Compute multiplication between tensor and scalar (field). @@ -251,7 +274,7 @@ def tensor_times_scalar( def apply_orthogonal_group( - orthogonal_matrix: IPhlowerTensor, tensor: IPhlowerTensor + orthogonal_matrix: IPhlowerTensor, tensor: IPhlowerTensor ) -> IPhlowerTensor: """ Apply orthogonal group action to the input tensor. @@ -277,16 +300,19 @@ def apply_orthogonal_group( start_dim = start_dim + 3 if tensor.is_voxel else start_dim s = _availale_variables(rank * 2) - str_ortho = ','.join(a + b for a, b in zip(s[::2], s[1::2], strict=True)) + str_ortho = ",".join(a + b for a, b in zip(s[::2], s[1::2], strict=True)) str_tensor = f"{time}{space}{s[1::2]}f" str_ret = f"{time}{space}{s[::2]}f" equation = f"{str_ortho},{str_tensor}->{str_ret}" args = [orthogonal_matrix] * rank + [tensor] return einsum( - equation, *args, + equation, + *args, dimension=tensor.dimension, - is_time_series=tensor.is_time_series, is_voxel=tensor.is_voxel) + is_time_series=tensor.is_time_series, + is_voxel=tensor.is_voxel, + ) def spatial_sum(tensor: IPhlowerTensor) -> IPhlowerTensor: @@ -298,9 +324,12 @@ def spatial_sum(tensor: IPhlowerTensor) -> IPhlowerTensor: space_width = 3 if tensor.is_voxel else 1 squeezed = einsum( - f"{time}{space}...->{time}...", tensor, - dimension=tensor.dimension, is_time_series=tensor.is_time_series, - is_voxel=tensor.is_voxel) + f"{time}{space}...->{time}...", + tensor, + dimension=tensor.dimension, + is_time_series=tensor.is_time_series, + is_voxel=tensor.is_voxel, + ) # keepdim for _ in range(space_width): diff --git a/src/phlower/nn/_core_modules/_pinv_mlp.py b/src/phlower/nn/_core_modules/_pinv_mlp.py index 69f9c41..8b4e950 100644 --- a/src/phlower/nn/_core_modules/_pinv_mlp.py +++ b/src/phlower/nn/_core_modules/_pinv_mlp.py @@ -59,7 +59,8 @@ def _initialize(self): """Initialize parameters and activations after resolve() is called.""" self._activations = [ _utils.ActivationSelector.select_inverse(name) - for name in self._reference._activations[::-1]] + for name in self._reference._activations[::-1] + ] self._chains = self._init_pinv_chains() def _init_pinv_chains(self): @@ -69,8 +70,7 @@ def _init_pinv_chains(self): raise ValueError(f"Unsupported reference class: {name}") - def _init_pinv_mlp_chains( - self, chains: _utils.ExtendedLinearList): + def _init_pinv_mlp_chains(self, chains: _utils.ExtendedLinearList): return [PInvLinear(c) for c in chains._linears[::-1]] def forward( @@ -100,7 +100,8 @@ def forward( h = data.unique_item() for activation, chain in zip( - self._activations, self._chains, strict=True): + self._activations, self._chains, strict=True + ): # Activation comes first because it is pseudo inverse h = chain(activation(h)) @@ -108,7 +109,6 @@ def forward( class PInvLinear(torch.nn.Module): - def __init__(self, ref_linear: torch.nn.Linear): super().__init__() self.ref_linear = ref_linear @@ -130,4 +130,4 @@ def bias(self): if self.ref_linear.bias is None: return 0 else: - return - self.ref_linear.bias + return -self.ref_linear.bias diff --git a/src/phlower/nn/_core_modules/_proportional.py b/src/phlower/nn/_core_modules/_proportional.py index 35e211e..bbcd0ca 100644 --- a/src/phlower/nn/_core_modules/_proportional.py +++ b/src/phlower/nn/_core_modules/_proportional.py @@ -52,8 +52,8 @@ def __init__( dropouts = [] self._chains = _utils.ExtendedLinearList( - nodes=nodes, activations=['identity'], - dropouts=[], bias=False) + nodes=nodes, activations=["identity"], dropouts=[], bias=False + ) self._nodes = nodes def resolve( diff --git a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py index 0de0035..750c943 100644 --- a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py @@ -81,14 +81,18 @@ def __init__( if self._disable_en_equivariance: self._mlp = MLP( - nodes=nodes, activations=activations, - dropouts=dropouts, bias=bias, + nodes=nodes, + activations=activations, + dropouts=dropouts, + bias=bias, ) self._linear_weight = self._init_linear_weight() else: self._mlp = EnEquivariantMLP( - nodes=nodes, activations=activations, - dropouts=dropouts, bias=bias, + nodes=nodes, + activations=activations, + dropouts=dropouts, + bias=bias, create_linear_weight=create_linear_weight, norm_function_name=norm_function_name, ) @@ -99,7 +103,8 @@ def _init_linear_weight(self): if self._nodes[0] != self._nodes[-1]: raise ValueError( "First and last nodes are different. " - "Set create_linear_weight True.") + "Set create_linear_weight True." + ) return Identity() return Proportional([self._nodes[0], self._nodes[-1]]) @@ -135,26 +140,33 @@ def forward( raise PhlowerDimensionRequiredError("Dimension is required") if dict_scales is None or len(data) < 1: raise PhlowerInvalidArgumentsError( - f"Scale inputs are required. Given: {dict_scales}, {kwards}") - if not np.all([ + f"Scale inputs are required. Given: {dict_scales}, {kwards}" + ) + if not np.all( + [ PhysicalDimensionSymbolType.is_exist(k) - for k in dict_scales.keys()]): + for k in dict_scales.keys() + ] + ): raise PhlowerInvalidArgumentsError( "keys in dict_scales should be in " - "PhysicalDimensionSymbolType. Given: {dict_scales.keys()}") + "PhysicalDimensionSymbolType. Given: {dict_scales.keys()}" + ) dict_dimension = h.dimension.to_dict() dict_scales = { - k: v**dict_dimension[k] for k, v in dict_scales.items()} + k: v ** dict_dimension[k] for k, v in dict_scales.items() + } # Make h dimensionless for v in dict_scales.values(): h = _functions.tensor_times_scalar(h, 1 / v) if self._centering: - volume = dict_dimension["L"]**3 + volume = dict_dimension["L"] ** 3 mean = _functions.spatial_mean( - _functions.tensor_times_scalar(h, volume)) + _functions.tensor_times_scalar(h, volume) + ) h = h - mean h = self._mlp(phlower_tensor_collection({"h": h})) @@ -162,7 +174,8 @@ def forward( if self._centering: linear_mean = self._linear_weight( - phlower_tensor_collection({"mean": mean})) + phlower_tensor_collection({"mean": mean}) + ) h = h + linear_mean if self._invariant: diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index 96afb8e..23b8b20 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -99,11 +99,13 @@ def select(name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: @staticmethod def select_inverse( - name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: + name: str | None, + ) -> Callable[[torch.Tensor], torch.Tensor]: if name is None: name = "identity" return ActivationSelector._REGISTERED_ACTIVATIONS[ - ActivationSelector._inverse_activation_name(name)] + ActivationSelector._inverse_activation_name(name) + ] @staticmethod def _inverse_activation_name(activation_name: str) -> str: @@ -117,7 +119,8 @@ def _inverse_activation_name(activation_name: str) -> str: return "truncated_atanh" raise PhlowerInvalidActivationError( - f"Cannot inverse for {activation_name}") + f"Cannot inverse for {activation_name}" + ) @staticmethod def is_exists(name: str) -> bool: diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index d3761fc..f8eb5aa 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -16,8 +16,9 @@ ProportionalSetting, ) from phlower.settings._module_settings._share_setting import ShareSetting -from phlower.settings._module_settings \ - ._similarity_equivariant_mlp_setting import SimilarityEquivariantMLPSetting +from phlower.settings._module_settings._similarity_equivariant_mlp_setting import ( + SimilarityEquivariantMLPSetting, +) _name_to_setting: dict[str, IPhlowerLayerParameters] = { "Concatenator": ConcatenatorSetting, diff --git a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py index 9774702..13069bf 100644 --- a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py @@ -19,7 +19,8 @@ class EnEquivariantMLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): bias: bool = Field(False, frozen=True) create_linear_weight: bool = Field(False, frozen=True) norm_function_name: str = Field( - default_factory=lambda: 'identity', frozen=True) + default_factory=lambda: "identity", frozen=True + ) def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: diff --git a/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py index 7c955e1..abc762e 100644 --- a/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py @@ -11,7 +11,8 @@ class SimilarityEquivariantMLPSetting( - IPhlowerLayerParameters, pydantic.BaseModel): + IPhlowerLayerParameters, pydantic.BaseModel +): nodes: list[int] = Field( ... ) # This property only overwritten when resolving. @@ -20,7 +21,8 @@ class SimilarityEquivariantMLPSetting( bias: bool = Field(False, frozen=True) create_linear_weight: bool = Field(False, frozen=True) norm_function_name: str = Field( - default_factory=lambda: 'identity', frozen=True) + default_factory=lambda: "identity", frozen=True + ) disable_en_equivariance: bool = Field(False, frozen=True) invariant: bool = Field(False, frozen=True) centering: bool = Field(False, frozen=True) @@ -28,7 +30,8 @@ class SimilarityEquivariantMLPSetting( def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError( - "Only one input is allowed in SimilarityEquivariantMLP.") + "Only one input is allowed in SimilarityEquivariantMLP." + ) return input_dims[0] @pydantic.field_validator("nodes") diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index 070e54a..a7cda1b 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -73,16 +73,20 @@ class PhlowerIncompatibleTensorError(ValueError): tensor(s) """ + class PhlowerTypeError(TypeError): """This error raises when type is not compatible with what is expected.""" + class PhlowerInvalidActivationError(ValueError): """This error raises when a set activation is invalid""" + class PhlowerDimensionRequiredError(ValueError): """ This error raises when the dimension does not exist despite required. """ + class PhlowerInvalidArgumentsError(ValueError): """This error raises when the arguments are invalid.""" diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index 675494c..e931701 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -111,10 +111,10 @@ def test__tensor_div_scalar(): dims = phlower_dimension_tensor({"L": 2, "T": -2}) a = torch.tensor(np.random.rand(3, 10)) - c = a / 3. + c = a / 3.0 ap = PhlowerTensor(a, dims) - cp = ap / 3. + cp = ap / 3.0 np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) @@ -126,10 +126,10 @@ def test__scalar_div_tensor(): desired_dims = phlower_dimension_tensor({"L": -2, "T": 2}) a = torch.tensor(np.random.rand(3, 10)) - c = 3. / a + c = 3.0 / a ap = PhlowerTensor(a, dims) - cp = 3. / ap + cp = 3.0 / ap np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) diff --git a/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py index ffae0f3..abdb50e 100644 --- a/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py +++ b/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py @@ -34,25 +34,32 @@ def test__can_call_parameters(): ], ) @pytest.mark.parametrize("n_output_feature", [1, 16, 32]) -def test__en_equivariance( - size, is_time_series, is_voxel, n_output_feature): +def test__en_equivariance(size, is_time_series, is_voxel, n_output_feature): orthogonal_tensor = PhlowerTensor( - torch.tensor(ortho_group.rvs(3).astype(np.float32))) + torch.tensor(ortho_group.rvs(3).astype(np.float32)) + ) create_linear_weight = size[-1] != n_output_feature model = EnEquivariantMLP( nodes=[size[-1], n_output_feature], - create_linear_weight=create_linear_weight) + create_linear_weight=create_linear_weight, + ) phlower_tensor = PhlowerTensor( - torch.rand(*size), is_time_series=is_time_series, is_voxel=is_voxel) + torch.rand(*size), is_time_series=is_time_series, is_voxel=is_voxel + ) - phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) + phlower_tensors = phlower_tensor_collection({"tensor": phlower_tensor}) actual = _functions.apply_orthogonal_group( - orthogonal_tensor, model(phlower_tensors)).to_numpy() + orthogonal_tensor, model(phlower_tensors) + ).to_numpy() rotated_phlower_tensors = phlower_tensor_collection( - {'tensor': _functions.apply_orthogonal_group( - orthogonal_tensor, phlower_tensor)}) + { + "tensor": _functions.apply_orthogonal_group( + orthogonal_tensor, phlower_tensor + ) + } + ) desired = model(rotated_phlower_tensors).to_numpy() np.testing.assert_almost_equal(actual, desired, decimal=6) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 8e8fb2b..3f4c6fd 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -49,7 +49,8 @@ def assert_correct(actual, array): desired = sp_sparse @ desired norm = np.mean(np.linalg.norm(desired, axis=-1)) np.testing.assert_almost_equal( - actual / norm, desired / norm, decimal=5) + actual / norm, desired / norm, decimal=5 + ) return for i in range(array.shape[1]): @@ -64,7 +65,7 @@ def assert_correct(actual, array): def test_leaky_relu0p5_inverse_leaky_relu0p5(): - x = PhlowerTensor(torch.rand(100)) * 4 - 2. + x = PhlowerTensor(torch.rand(100)) * 4 - 2.0 y = _functions.inversed_leaky_relu0p5(_functions.leaky_relu0p5(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) @@ -113,14 +114,17 @@ def test_smooth_leaky_relu_inverse(): ], ) def test_contraction_one_argument( - size, is_time_series, is_voxel, desired_pattern, dimension): + size, is_time_series, is_voxel, desired_pattern, dimension +): torch_tensor = torch.rand(*size) x = phlower_tensor( - torch_tensor, dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + torch_tensor, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) actual = _functions.contraction(x) - desired = torch.einsum( - desired_pattern, torch_tensor, torch_tensor).numpy() + desired = torch.einsum(desired_pattern, torch_tensor, torch_tensor).numpy() np.testing.assert_almost_equal(actual.to_tensor().numpy(), desired) assert actual.is_time_series == is_time_series @@ -140,137 +144,237 @@ def test_contraction_one_argument( "desired_pattern, desired_rank", [ # Base - ( - (10, 16), (10, 16), False, False, False, - "nf,nf->nf", 0), - ( - (10, 3, 16), (10, 16), False, False, False, - "npf,nf->npf", 1), - ( - (10, 16), (10, 3, 16), False, False, False, - "nf,npf->npf", 1), - ( - (10, 3, 16), (10, 3, 16), False, False, False, - "npf,npf->nf", 0), - ( - (10, 3, 3, 16), (10, 16), False, False, False, - "npqf,nf->npqf", 2), - ( - (10, 3, 3, 16), (10, 3, 16), False, False, False, - "npqf,npf->nqf", 1), - ( - (10, 3, 3, 16), (10, 3, 3, 16), False, False, False, - "npqf,npqf->nf", 0), + ((10, 16), (10, 16), False, False, False, "nf,nf->nf", 0), + ((10, 3, 16), (10, 16), False, False, False, "npf,nf->npf", 1), + ((10, 16), (10, 3, 16), False, False, False, "nf,npf->npf", 1), + ((10, 3, 16), (10, 3, 16), False, False, False, "npf,npf->nf", 0), + ((10, 3, 3, 16), (10, 16), False, False, False, "npqf,nf->npqf", 2), + ((10, 3, 3, 16), (10, 3, 16), False, False, False, "npqf,npf->nqf", 1), + ( + (10, 3, 3, 16), + (10, 3, 3, 16), + False, + False, + False, + "npqf,npqf->nf", + 0, + ), # X time series - ( - (4, 10, 16), (10, 16), True, False, False, - "tnf,nf->tnf", 0), - ( - (4, 10, 3, 16), (10, 16), True, False, False, - "tnpf,nf->tnpf", 1), - ( - (4, 10, 16), (10, 3, 16), True, False, False, - "tnf,npf->tnpf", 1), - ( - (4, 10, 3, 16), (10, 3, 16), True, False, False, - "tnpf,npf->tnf", 0), - ( - (4, 10, 3, 3, 16), (10, 16), True, False, False, - "tnpqf,nf->tnpqf", 2), - ( - (4, 10, 3, 3, 16), (10, 3, 16), True, False, False, - "tnpqf,npf->tnqf", 1), - ( - (4, 10, 3, 3, 16), (10, 3, 3, 16), True, False, False, - "tnpqf,npqf->tnf", 0), + ((4, 10, 16), (10, 16), True, False, False, "tnf,nf->tnf", 0), + ((4, 10, 3, 16), (10, 16), True, False, False, "tnpf,nf->tnpf", 1), + ((4, 10, 16), (10, 3, 16), True, False, False, "tnf,npf->tnpf", 1), + ((4, 10, 3, 16), (10, 3, 16), True, False, False, "tnpf,npf->tnf", 0), + ((4, 10, 3, 3, 16), (10, 16), True, False, False, "tnpqf,nf->tnpqf", 2), + ( + (4, 10, 3, 3, 16), + (10, 3, 16), + True, + False, + False, + "tnpqf,npf->tnqf", + 1, + ), + ( + (4, 10, 3, 3, 16), + (10, 3, 3, 16), + True, + False, + False, + "tnpqf,npqf->tnf", + 0, + ), # Y time series - ( - (10, 16), (4, 10, 16), False, True, False, - "nf,tnf->tnf", 0), - ( - (10, 3, 16), (4, 10, 16), False, True, False, - "npf,tnf->tnpf", 1), - ( - (10, 16), (4, 10, 3, 16), False, True, False, - "nf,tnpf->tnpf", 1), - ( - (10, 3, 16), (4, 10, 3, 16), False, True, False, - "npf,tnpf->tnf", 0), - ( - (10, 3, 3, 16), (4, 10, 16), False, True, False, - "npqf,tnf->tnpqf", 2), - ( - (10, 3, 3, 16), (4, 10, 3, 16), False, True, False, - "npqf,tnpf->tnqf", 1), - ( - (10, 3, 3, 16), (4, 10, 3, 3, 16), False, True, False, - "npqf,tnpqf->tnf", 0), + ((10, 16), (4, 10, 16), False, True, False, "nf,tnf->tnf", 0), + ((10, 3, 16), (4, 10, 16), False, True, False, "npf,tnf->tnpf", 1), + ((10, 16), (4, 10, 3, 16), False, True, False, "nf,tnpf->tnpf", 1), + ((10, 3, 16), (4, 10, 3, 16), False, True, False, "npf,tnpf->tnf", 0), + ((10, 3, 3, 16), (4, 10, 16), False, True, False, "npqf,tnf->tnpqf", 2), + ( + (10, 3, 3, 16), + (4, 10, 3, 16), + False, + True, + False, + "npqf,tnpf->tnqf", + 1, + ), + ( + (10, 3, 3, 16), + (4, 10, 3, 3, 16), + False, + True, + False, + "npqf,tnpqf->tnf", + 0, + ), # X Y time series - ( - (4, 10, 16), (4, 10, 16), True, True, False, - "tnf,tnf->tnf", 0), - ( - (4, 10, 3, 16), (4, 10, 16), True, True, False, - "tnpf,tnf->tnpf", 1), - ( - (4, 10, 16), (4, 10, 3, 16), True, True, False, - "tnf,tnpf->tnpf", 1), - ( - (4, 10, 3, 16), (4, 10, 3, 16), True, True, False, - "tnpf,tnpf->tnf", 0), - ( - (4, 10, 3, 3, 16), (4, 10, 16), True, True, False, - "tnpqf,tnf->tnpqf", 2), - ( - (4, 10, 3, 3, 16), (4, 10, 3, 16), True, True, False, - "tnpqf,tnpf->tnqf", 1), - ( - (4, 10, 3, 3, 16), (4, 10, 3, 3, 16), True, True, False, - "tnpqf,tnpqf->tnf", 0), + ((4, 10, 16), (4, 10, 16), True, True, False, "tnf,tnf->tnf", 0), + ((4, 10, 3, 16), (4, 10, 16), True, True, False, "tnpf,tnf->tnpf", 1), + ((4, 10, 16), (4, 10, 3, 16), True, True, False, "tnf,tnpf->tnpf", 1), + ( + (4, 10, 3, 16), + (4, 10, 3, 16), + True, + True, + False, + "tnpf,tnpf->tnf", + 0, + ), + ( + (4, 10, 3, 3, 16), + (4, 10, 16), + True, + True, + False, + "tnpqf,tnf->tnpqf", + 2, + ), + ( + (4, 10, 3, 3, 16), + (4, 10, 3, 16), + True, + True, + False, + "tnpqf,tnpf->tnqf", + 1, + ), + ( + (4, 10, 3, 3, 16), + (4, 10, 3, 3, 16), + True, + True, + False, + "tnpqf,tnpqf->tnf", + 0, + ), # X time series, X Y voxel ( - (4, 10, 10, 10, 16), (10, 10, 10, 16), - True, False, True, "txyzf,xyzf->txyzf", 0), - ( - (4, 10, 10, 10, 3, 16), (10, 10, 10, 16), - True, False, True, "txyzpf,xyzf->txyzpf", 1), - ( - (4, 10, 10, 10, 16), (10, 10, 10, 3, 16), - True, False, True, "txyzf,xyzpf->txyzpf", 1), - ( - (4, 10, 10, 10, 3, 16), (10, 10, 10, 3, 16), - True, False, True, "txyzpf,xyzpf->txyzf", 0), - ( - (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 16), - True, False, True, "txyzpqf,xyzf->txyzpqf", 2), - ( - (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 16), - True, False, True, "txyzpqf,xyzpf->txyzqf", 1), - ( - (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 3, 16), - True, False, True, "txyzpqf,xyzpqf->txyzf", 0), + (4, 10, 10, 10, 16), + (10, 10, 10, 16), + True, + False, + True, + "txyzf,xyzf->txyzf", + 0, + ), + ( + (4, 10, 10, 10, 3, 16), + (10, 10, 10, 16), + True, + False, + True, + "txyzpf,xyzf->txyzpf", + 1, + ), + ( + (4, 10, 10, 10, 16), + (10, 10, 10, 3, 16), + True, + False, + True, + "txyzf,xyzpf->txyzpf", + 1, + ), + ( + (4, 10, 10, 10, 3, 16), + (10, 10, 10, 3, 16), + True, + False, + True, + "txyzpf,xyzpf->txyzf", + 0, + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (10, 10, 10, 16), + True, + False, + True, + "txyzpqf,xyzf->txyzpqf", + 2, + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (10, 10, 10, 3, 16), + True, + False, + True, + "txyzpqf,xyzpf->txyzqf", + 1, + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (10, 10, 10, 3, 3, 16), + True, + False, + True, + "txyzpqf,xyzpqf->txyzf", + 0, + ), # X Y time series, X Y voxel ( - (4, 10, 10, 10, 16), (4, 10, 10, 10, 16), - True, True, True, "txyzf,txyzf->txyzf", 0), - ( - (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 16), - True, True, True, "txyzpf,txyzf->txyzpf", 1), - ( - (4, 10, 10, 10, 16), (4, 10, 10, 10, 3, 16), - True, True, True, "txyzf,txyzpf->txyzpf", 1), - ( - (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 3, 16), - True, True, True, "txyzpf,txyzpf->txyzf", 0), - ( - (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 16), - True, True, True, "txyzpqf,txyzf->txyzpqf", 2), - ( - (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 16), - True, True, True, "txyzpqf,txyzpf->txyzqf", 1), - ( - (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 3, 16), - True, True, True, "txyzpqf,txyzpqf->txyzf", 0), + (4, 10, 10, 10, 16), + (4, 10, 10, 10, 16), + True, + True, + True, + "txyzf,txyzf->txyzf", + 0, + ), + ( + (4, 10, 10, 10, 3, 16), + (4, 10, 10, 10, 16), + True, + True, + True, + "txyzpf,txyzf->txyzpf", + 1, + ), + ( + (4, 10, 10, 10, 16), + (4, 10, 10, 10, 3, 16), + True, + True, + True, + "txyzf,txyzpf->txyzpf", + 1, + ), + ( + (4, 10, 10, 10, 3, 16), + (4, 10, 10, 10, 3, 16), + True, + True, + True, + "txyzpf,txyzpf->txyzf", + 0, + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (4, 10, 10, 10, 16), + True, + True, + True, + "txyzpqf,txyzf->txyzpqf", + 2, + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (4, 10, 10, 10, 3, 16), + True, + True, + True, + "txyzpqf,txyzpf->txyzqf", + 1, + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (4, 10, 10, 10, 3, 3, 16), + True, + True, + True, + "txyzpqf,txyzpqf->txyzf", + 0, + ), ], ) @pytest.mark.parametrize( @@ -292,21 +396,34 @@ def test_contraction_one_argument( ], ) def test_contraction_two_arguments( - size_x, size_y, x_is_time_series, y_is_time_series, is_voxel, - desired_pattern, desired_rank, dimension_x, dimension_y): + size_x, + size_y, + x_is_time_series, + y_is_time_series, + is_voxel, + desired_pattern, + desired_rank, + dimension_x, + dimension_y, +): t_x = torch.rand(*size_x) x = phlower_tensor( - t_x, dimension=dimension_x, - is_time_series=x_is_time_series, is_voxel=is_voxel) + t_x, + dimension=dimension_x, + is_time_series=x_is_time_series, + is_voxel=is_voxel, + ) t_y = torch.rand(*size_y) y = phlower_tensor( - t_y, dimension=dimension_y, - is_time_series=y_is_time_series, is_voxel=is_voxel) + t_y, + dimension=dimension_y, + is_time_series=y_is_time_series, + is_voxel=is_voxel, + ) actual = _functions.contraction(x, y) - desired = torch.einsum( - desired_pattern, t_x, t_y).numpy() + desired = torch.einsum(desired_pattern, t_x, t_y).numpy() np.testing.assert_almost_equal(actual.to_tensor().numpy(), desired) assert actual.is_time_series == x_is_time_series or y_is_time_series @@ -322,10 +439,8 @@ def test_contraction_two_arguments( def test_contraction_raises_phlower_incompatible_tensor_error(): - x = phlower_tensor( - torch.rand(10, 10, 10, 3, 16), is_voxel=True) - y = phlower_tensor( - torch.rand(10 * 10 * 10, 3, 16), is_voxel=False) + x = phlower_tensor(torch.rand(10, 10, 10, 3, 16), is_voxel=True) + y = phlower_tensor(torch.rand(10 * 10 * 10, 3, 16), is_voxel=False) with pytest.raises(PhlowerIncompatibleTensorError): _functions.contraction(x, y) @@ -335,137 +450,200 @@ def test_contraction_raises_phlower_incompatible_tensor_error(): "desired_pattern", [ # Base - ( - (10, 16), (10, 16), False, False, False, - "nf,nf->nf"), - ( - (10, 3, 16), (10, 16), False, False, False, - "npf,nf->npf"), - ( - (10, 16), (10, 3, 16), False, False, False, - "nf,npf->npf"), - ( - (10, 3, 16), (10, 3, 16), False, False, False, - "npf,nqf->npqf"), - ( - (10, 3, 3, 16), (10, 16), False, False, False, - "npqf,nf->npqf"), - ( - (10, 3, 3, 16), (10, 3, 16), False, False, False, - "npqf,nrf->npqrf"), - ( - (10, 3, 3, 16), (10, 3, 3, 16), False, False, False, - "npqf,nrsf->npqrsf"), + ((10, 16), (10, 16), False, False, False, "nf,nf->nf"), + ((10, 3, 16), (10, 16), False, False, False, "npf,nf->npf"), + ((10, 16), (10, 3, 16), False, False, False, "nf,npf->npf"), + ((10, 3, 16), (10, 3, 16), False, False, False, "npf,nqf->npqf"), + ((10, 3, 3, 16), (10, 16), False, False, False, "npqf,nf->npqf"), + ((10, 3, 3, 16), (10, 3, 16), False, False, False, "npqf,nrf->npqrf"), + ( + (10, 3, 3, 16), + (10, 3, 3, 16), + False, + False, + False, + "npqf,nrsf->npqrsf", + ), # X time series - ( - (4, 10, 16), (10, 16), True, False, False, - "tnf,nf->tnf"), - ( - (4, 10, 3, 16), (10, 16), True, False, False, - "tnpf,nf->tnpf"), - ( - (4, 10, 16), (10, 3, 16), True, False, False, - "tnf,npf->tnpf"), - ( - (4, 10, 3, 16), (10, 3, 16), True, False, False, - "tnpf,nqf->tnpqf"), - ( - (4, 10, 3, 3, 16), (10, 16), True, False, False, - "tnpqf,nf->tnpqf"), - ( - (4, 10, 3, 3, 16), (10, 3, 16), True, False, False, - "tnpqf,nrf->tnpqrf"), - ( - (4, 10, 3, 3, 16), (10, 3, 3, 16), True, False, False, - "tnpqf,nrsf->tnpqrsf"), + ((4, 10, 16), (10, 16), True, False, False, "tnf,nf->tnf"), + ((4, 10, 3, 16), (10, 16), True, False, False, "tnpf,nf->tnpf"), + ((4, 10, 16), (10, 3, 16), True, False, False, "tnf,npf->tnpf"), + ((4, 10, 3, 16), (10, 3, 16), True, False, False, "tnpf,nqf->tnpqf"), + ((4, 10, 3, 3, 16), (10, 16), True, False, False, "tnpqf,nf->tnpqf"), + ( + (4, 10, 3, 3, 16), + (10, 3, 16), + True, + False, + False, + "tnpqf,nrf->tnpqrf", + ), + ( + (4, 10, 3, 3, 16), + (10, 3, 3, 16), + True, + False, + False, + "tnpqf,nrsf->tnpqrsf", + ), # Y time series - ( - (10, 16), (4, 10, 16), False, True, False, - "nf,tnf->tnf"), - ( - (10, 3, 16), (4, 10, 16), False, True, False, - "npf,tnf->tnpf"), - ( - (10, 16), (4, 10, 3, 16), False, True, False, - "nf,tnpf->tnpf"), - ( - (10, 3, 16), (4, 10, 3, 16), False, True, False, - "npf,tnqf->tnpqf"), - ( - (10, 3, 3, 16), (4, 10, 16), False, True, False, - "npqf,tnf->tnpqf"), - ( - (10, 3, 3, 16), (4, 10, 3, 16), False, True, False, - "npqf,tnrf->tnpqrf"), - ( - (10, 3, 3, 16), (4, 10, 3, 3, 16), False, True, False, - "npqf,tnrsf->tnpqrsf"), + ((10, 16), (4, 10, 16), False, True, False, "nf,tnf->tnf"), + ((10, 3, 16), (4, 10, 16), False, True, False, "npf,tnf->tnpf"), + ((10, 16), (4, 10, 3, 16), False, True, False, "nf,tnpf->tnpf"), + ((10, 3, 16), (4, 10, 3, 16), False, True, False, "npf,tnqf->tnpqf"), + ((10, 3, 3, 16), (4, 10, 16), False, True, False, "npqf,tnf->tnpqf"), + ( + (10, 3, 3, 16), + (4, 10, 3, 16), + False, + True, + False, + "npqf,tnrf->tnpqrf", + ), + ( + (10, 3, 3, 16), + (4, 10, 3, 3, 16), + False, + True, + False, + "npqf,tnrsf->tnpqrsf", + ), # X Y time series - ( - (4, 10, 16), (4, 10, 16), True, True, False, - "tnf,tnf->tnf"), - ( - (4, 10, 3, 16), (4, 10, 16), True, True, False, - "tnpf,tnf->tnpf"), - ( - (4, 10, 16), (4, 10, 3, 16), True, True, False, - "tnf,tnpf->tnpf"), - ( - (4, 10, 3, 16), (4, 10, 3, 16), True, True, False, - "tnpf,tnqf->tnpqf"), - ( - (4, 10, 3, 3, 16), (4, 10, 16), True, True, False, - "tnpqf,tnf->tnpqf"), - ( - (4, 10, 3, 3, 16), (4, 10, 3, 16), True, True, False, - "tnpqf,tnrf->tnpqrf"), - ( - (4, 10, 3, 3, 16), (4, 10, 3, 3, 16), True, True, False, - "tnpqf,tnrsf->tnpqrsf"), + ((4, 10, 16), (4, 10, 16), True, True, False, "tnf,tnf->tnf"), + ((4, 10, 3, 16), (4, 10, 16), True, True, False, "tnpf,tnf->tnpf"), + ((4, 10, 16), (4, 10, 3, 16), True, True, False, "tnf,tnpf->tnpf"), + ((4, 10, 3, 16), (4, 10, 3, 16), True, True, False, "tnpf,tnqf->tnpqf"), + ((4, 10, 3, 3, 16), (4, 10, 16), True, True, False, "tnpqf,tnf->tnpqf"), + ( + (4, 10, 3, 3, 16), + (4, 10, 3, 16), + True, + True, + False, + "tnpqf,tnrf->tnpqrf", + ), + ( + (4, 10, 3, 3, 16), + (4, 10, 3, 3, 16), + True, + True, + False, + "tnpqf,tnrsf->tnpqrsf", + ), # X time series, X Y voxel ( - (4, 10, 10, 10, 16), (10, 10, 10, 16), - True, False, True, "txyzf,xyzf->txyzf"), - ( - (4, 10, 10, 10, 3, 16), (10, 10, 10, 16), - True, False, True, "txyzpf,xyzf->txyzpf"), - ( - (4, 10, 10, 10, 16), (10, 10, 10, 3, 16), - True, False, True, "txyzf,xyzpf->txyzpf"), - ( - (4, 10, 10, 10, 3, 16), (10, 10, 10, 3, 16), - True, False, True, "txyzpf,xyzqf->txyzpqf"), - ( - (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 16), - True, False, True, "txyzpqf,xyzf->txyzpqf"), - ( - (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 16), - True, False, True, "txyzpqf,xyzrf->txyzpqrf"), - ( - (4, 10, 10, 10, 3, 3, 16), (10, 10, 10, 3, 3, 16), - True, False, True, "txyzpqf,xyzrsf->txyzpqrsf"), + (4, 10, 10, 10, 16), + (10, 10, 10, 16), + True, + False, + True, + "txyzf,xyzf->txyzf", + ), + ( + (4, 10, 10, 10, 3, 16), + (10, 10, 10, 16), + True, + False, + True, + "txyzpf,xyzf->txyzpf", + ), + ( + (4, 10, 10, 10, 16), + (10, 10, 10, 3, 16), + True, + False, + True, + "txyzf,xyzpf->txyzpf", + ), + ( + (4, 10, 10, 10, 3, 16), + (10, 10, 10, 3, 16), + True, + False, + True, + "txyzpf,xyzqf->txyzpqf", + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (10, 10, 10, 16), + True, + False, + True, + "txyzpqf,xyzf->txyzpqf", + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (10, 10, 10, 3, 16), + True, + False, + True, + "txyzpqf,xyzrf->txyzpqrf", + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (10, 10, 10, 3, 3, 16), + True, + False, + True, + "txyzpqf,xyzrsf->txyzpqrsf", + ), # X Y time series, X Y voxel ( - (4, 10, 10, 10, 16), (4, 10, 10, 10, 16), - True, True, True, "txyzf,txyzf->txyzf"), - ( - (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 16), - True, True, True, "txyzpf,txyzf->txyzpf"), - ( - (4, 10, 10, 10, 16), (4, 10, 10, 10, 3, 16), - True, True, True, "txyzf,txyzpf->txyzpf"), - ( - (4, 10, 10, 10, 3, 16), (4, 10, 10, 10, 3, 16), - True, True, True, "txyzpf,txyzqf->txyzpqf"), - ( - (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 16), - True, True, True, "txyzpqf,txyzf->txyzpqf"), - ( - (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 16), - True, True, True, "txyzpqf,txyzrf->txyzpqrf"), - ( - (4, 10, 10, 10, 3, 3, 16), (4, 10, 10, 10, 3, 3, 16), - True, True, True, "txyzpqf,txyzrsf->txyzpqrsf"), + (4, 10, 10, 10, 16), + (4, 10, 10, 10, 16), + True, + True, + True, + "txyzf,txyzf->txyzf", + ), + ( + (4, 10, 10, 10, 3, 16), + (4, 10, 10, 10, 16), + True, + True, + True, + "txyzpf,txyzf->txyzpf", + ), + ( + (4, 10, 10, 10, 16), + (4, 10, 10, 10, 3, 16), + True, + True, + True, + "txyzf,txyzpf->txyzpf", + ), + ( + (4, 10, 10, 10, 3, 16), + (4, 10, 10, 10, 3, 16), + True, + True, + True, + "txyzpf,txyzqf->txyzpqf", + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (4, 10, 10, 10, 16), + True, + True, + True, + "txyzpqf,txyzf->txyzpqf", + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (4, 10, 10, 10, 3, 16), + True, + True, + True, + "txyzpqf,txyzrf->txyzpqrf", + ), + ( + (4, 10, 10, 10, 3, 3, 16), + (4, 10, 10, 10, 3, 3, 16), + True, + True, + True, + "txyzpqf,txyzrsf->txyzpqrsf", + ), ], ) @pytest.mark.parametrize( @@ -487,21 +665,33 @@ def test_contraction_raises_phlower_incompatible_tensor_error(): ], ) def test_tensor_product( - size_x, size_y, x_is_time_series, y_is_time_series, is_voxel, - desired_pattern, dimension_x, dimension_y): + size_x, + size_y, + x_is_time_series, + y_is_time_series, + is_voxel, + desired_pattern, + dimension_x, + dimension_y, +): t_x = torch.rand(*size_x) x = phlower_tensor( - t_x, dimension=dimension_x, - is_time_series=x_is_time_series, is_voxel=is_voxel) + t_x, + dimension=dimension_x, + is_time_series=x_is_time_series, + is_voxel=is_voxel, + ) t_y = torch.rand(*size_y) y = phlower_tensor( - t_y, dimension=dimension_y, - is_time_series=y_is_time_series, is_voxel=is_voxel) + t_y, + dimension=dimension_y, + is_time_series=y_is_time_series, + is_voxel=is_voxel, + ) actual = _functions.tensor_product(x, y) - desired = torch.einsum( - desired_pattern, t_x, t_y).numpy() + desired = torch.einsum(desired_pattern, t_x, t_y).numpy() np.testing.assert_almost_equal(actual.to_numpy(), desired) assert actual.is_time_series == x_is_time_series or y_is_time_series @@ -547,22 +737,25 @@ def test_tensor_product( ], ) def test_apply_orthogonal_group( - size, is_time_series, is_voxel, desired_pattern, dimension): + size, is_time_series, is_voxel, desired_pattern, dimension +): orthogonal_matrix = torch.from_numpy(ortho_group.rvs(3).astype(np.float32)) orthogonal_tensor = phlower_tensor(orthogonal_matrix) torch_tensor = torch.rand(*size) x = phlower_tensor( - torch_tensor, dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + torch_tensor, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) actual = _functions.apply_orthogonal_group(orthogonal_tensor, x) if desired_pattern is None: desired = torch_tensor.numpy() else: inputs = [orthogonal_matrix] * x.rank() + [torch_tensor] - desired = torch.einsum( - desired_pattern, *inputs).numpy() + desired = torch.einsum(desired_pattern, *inputs).numpy() np.testing.assert_almost_equal(actual.to_numpy(), desired) assert actual.is_time_series == is_time_series @@ -607,19 +800,22 @@ def test_apply_orthogonal_group( [[-1], [-1], [2], [0], [1], [0], [0]], ], ) -def test_spatial_mean( - size, is_time_series, is_voxel, mean_dims, dimension): +def test_spatial_mean(size, is_time_series, is_voxel, mean_dims, dimension): torch_tensor = torch.rand(*size) x = phlower_tensor( - torch_tensor, dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + torch_tensor, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) actual = _functions.spatial_mean(x) desired = torch_tensor for dim in mean_dims: desired = torch.mean(desired, dim=dim, keepdim=True) np.testing.assert_almost_equal( - actual.to_numpy(), desired.numpy(), decimal=5) + actual.to_numpy(), desired.numpy(), decimal=5 + ) assert actual.is_time_series == is_time_series assert actual.is_voxel == is_voxel diff --git a/tests/test_nn/test_core_modules/test_identity.py b/tests/test_nn/test_core_modules/test_identity.py index 977d260..5e07b1e 100644 --- a/tests/test_nn/test_core_modules/test_identity.py +++ b/tests/test_nn/test_core_modules/test_identity.py @@ -24,11 +24,11 @@ def test__can_call_parameters(): def test__identity(input_shape): phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) phlower_tensors = phlower_tensor_collection( - {"phlower_tensor": phlower_tensor}) + {"phlower_tensor": phlower_tensor} + ) model = Identity() actual = model(phlower_tensors) - np.testing.assert_almost_equal( - actual.to_numpy(), phlower_tensor.to_numpy()) + np.testing.assert_almost_equal(actual.to_numpy(), phlower_tensor.to_numpy()) diff --git a/tests/test_nn/test_core_modules/test_pinv.py b/tests/test_nn/test_core_modules/test_pinv.py index bd8b25b..cf7bdde 100644 --- a/tests/test_nn/test_core_modules/test_pinv.py +++ b/tests/test_nn/test_core_modules/test_pinv.py @@ -1,4 +1,3 @@ - import numpy as np import pytest import torch @@ -40,4 +39,5 @@ def test__pinv_mlp(mlp_nodes, activations, decimal): pinv_val = model(phlower_tensor_collection({"tensor": mlp_val})) np.testing.assert_array_almost_equal( - pinv_val.to_numpy(), t.to_numpy(), decimal=decimal) + pinv_val.to_numpy(), t.to_numpy(), decimal=decimal + ) diff --git a/tests/test_nn/test_core_modules/test_proportional.py b/tests/test_nn/test_core_modules/test_proportional.py index f7cc7bd..3b51f05 100644 --- a/tests/test_nn/test_core_modules/test_proportional.py +++ b/tests/test_nn/test_core_modules/test_proportional.py @@ -26,20 +26,20 @@ def test__can_call_parameters(): ], ) @pytest.mark.parametrize("n_output_feature", [1, 16, 32]) -@pytest.mark.parametrize("scale", [0., 0.5, 2.]) -def test__proportional_linearity( - size, is_time_series, n_output_feature, scale): - +@pytest.mark.parametrize("scale", [0.0, 0.5, 2.0]) +def test__proportional_linearity(size, is_time_series, n_output_feature, scale): model = Proportional(nodes=[size[-1], n_output_feature]) phlower_tensor = PhlowerTensor( - torch.rand(*size), is_time_series=is_time_series) + torch.rand(*size), is_time_series=is_time_series + ) - phlower_tensors = phlower_tensor_collection({'tensor': phlower_tensor}) + phlower_tensors = phlower_tensor_collection({"tensor": phlower_tensor}) actual = model(phlower_tensors).to_numpy() scaled_phlower_tensors = phlower_tensor_collection( - {'tensor': phlower_tensor * scale}) + {"tensor": phlower_tensor * scale} + ) desired = model(scaled_phlower_tensors).to_numpy() np.testing.assert_almost_equal(actual * scale, desired) diff --git a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py index 603144f..5cb9b59 100644 --- a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py +++ b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py @@ -48,26 +48,32 @@ def test__can_call_parameters(): # Velocity {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, # Mass density - {"T": 0, "L": - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + {"T": 0, "L": -3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, # Momentum density {"T": -1, "L": 1 - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, ], ) -@pytest.mark.parametrize( - "norm_function_name", ["identity", "sqrt"]) -@pytest.mark.parametrize( - "centering", [False, True]) +@pytest.mark.parametrize("norm_function_name", ["identity", "sqrt"]) +@pytest.mark.parametrize("centering", [False, True]) def test__similarity_equivariance( - size, is_time_series, is_voxel, activation, n_output_feature, - dimension, norm_function_name, centering, + size, + is_time_series, + is_voxel, + activation, + n_output_feature, + dimension, + norm_function_name, + centering, ): orthogonal_tensor = PhlowerTensor( - torch.tensor(ortho_group.rvs(3).astype(np.float32))) + torch.tensor(ortho_group.rvs(3).astype(np.float32)) + ) dict_scaling_factor = { - k: - np.random.rand() * 10 for k, v in dimension.items()} + k: np.random.rand() * 10 for k, v in dimension.items() + } scaling_factor = np.prod( - [dict_scaling_factor[k]**v for k, v in dimension.items()]) + [dict_scaling_factor[k] ** v for k, v in dimension.items()] + ) create_linear_weight = size[-1] != n_output_feature model = SimilarityEquivariantMLP( @@ -80,29 +86,40 @@ def test__similarity_equivariance( ) t = phlower_tensor( - torch.rand(*size) * 2 - 1., dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + torch.rand(*size) * 2 - 1.0, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) dict_tensor = {"tensor": t} dict_scales = { k: phlower_tensor( - torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) + torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1}) + ) for k, v in dimension.items() } ts = phlower_tensor_collection(dict_tensor) - actual_tensor = _functions.apply_orthogonal_group( - orthogonal_tensor, - model(ts, dict_scales=dict_scales)) * scaling_factor + actual_tensor = ( + _functions.apply_orthogonal_group( + orthogonal_tensor, model(ts, dict_scales=dict_scales) + ) + * scaling_factor + ) actual = actual_tensor.to_numpy() - dict_transformed = {'tensor': _functions.apply_orthogonal_group( - orthogonal_tensor, t) * scaling_factor} + dict_transformed = { + "tensor": _functions.apply_orthogonal_group(orthogonal_tensor, t) + * scaling_factor + } dict_scaled_scales = { - k: v * dict_scaling_factor[k] for k, v in dict_scales.items()} + k: v * dict_scaling_factor[k] for k, v in dict_scales.items() + } transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) desired = model( - transformed_phlower_tensors, dict_scales=dict_scaled_scales).to_numpy() + transformed_phlower_tensors, dict_scales=dict_scaled_scales + ).to_numpy() scale = np.max(np.abs(desired)) # Test equivariance @@ -140,26 +157,32 @@ def test__similarity_equivariance( # Velocity {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, # Mass density - {"T": 0, "L": - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + {"T": 0, "L": -3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, # Momentum density {"T": -1, "L": 1 - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, ], ) -@pytest.mark.parametrize( - "norm_function_name", ["identity", "sqrt"]) -@pytest.mark.parametrize( - "centering", [False, True]) +@pytest.mark.parametrize("norm_function_name", ["identity", "sqrt"]) +@pytest.mark.parametrize("centering", [False, True]) def test__similarity_invariance( - size, is_time_series, is_voxel, activation, n_output_feature, - dimension, norm_function_name, centering, + size, + is_time_series, + is_voxel, + activation, + n_output_feature, + dimension, + norm_function_name, + centering, ): orthogonal_tensor = PhlowerTensor( - torch.tensor(ortho_group.rvs(3).astype(np.float32))) + torch.tensor(ortho_group.rvs(3).astype(np.float32)) + ) dict_scaling_factor = { - k: - np.random.rand() * 10 for k, v in dimension.items()} + k: np.random.rand() * 10 for k, v in dimension.items() + } scaling_factor = np.prod( - [dict_scaling_factor[k]**v for k, v in dimension.items()]) + [dict_scaling_factor[k] ** v for k, v in dimension.items()] + ) create_linear_weight = size[-1] != n_output_feature model = SimilarityEquivariantMLP( @@ -168,32 +191,42 @@ def test__similarity_invariance( create_linear_weight=create_linear_weight, norm_function_name=norm_function_name, disable_en_equivariance=False, - centering=centering, invariant=True, + centering=centering, + invariant=True, ) t = phlower_tensor( - torch.rand(*size) * 2 - 1., dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + torch.rand(*size) * 2 - 1.0, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) dict_tensor = {"tensor": t} dict_scales = { k: phlower_tensor( - torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) + torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1}) + ) for k, v in dimension.items() } ts = phlower_tensor_collection(dict_tensor) actual_tensor = _functions.apply_orthogonal_group( - orthogonal_tensor, model(ts, dict_scales=dict_scales)) + orthogonal_tensor, model(ts, dict_scales=dict_scales) + ) actual = actual_tensor.to_numpy() - dict_transformed = {'tensor': _functions.apply_orthogonal_group( - orthogonal_tensor, t) * scaling_factor} + dict_transformed = { + "tensor": _functions.apply_orthogonal_group(orthogonal_tensor, t) + * scaling_factor + } dict_scaled_scales = { - k: v * dict_scaling_factor[k] for k, v in dict_scales.items()} + k: v * dict_scaling_factor[k] for k, v in dict_scales.items() + } transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) desired = model( - transformed_phlower_tensors, dict_scales=dict_scaled_scales).to_numpy() + transformed_phlower_tensors, dict_scales=dict_scaled_scales + ).to_numpy() scale = np.max(np.abs(desired)) # Test equivariance @@ -201,7 +234,7 @@ def test__similarity_invariance( # Test dimensionless in case of invariant for v in actual_tensor.dimension.to_dict().values(): - np.testing.assert_almost_equal(v, 0.) + np.testing.assert_almost_equal(v, 0.0) @pytest.mark.parametrize( @@ -231,28 +264,32 @@ def test__similarity_invariance( # Velocity {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0}, # Mass density - {"T": 0, "L": - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, + {"T": 0, "L": -3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, # Momentum density {"T": -1, "L": 1 - 3, "M": 1, "I": 0, "Theta": 0, "N": 0, "J": 0}, ], ) -@pytest.mark.parametrize( - "norm_function_name", ["identity", "sqrt"]) -@pytest.mark.parametrize( - "centering", [False, True]) -@pytest.mark.parametrize( - "invariant", [False, True]) +@pytest.mark.parametrize("norm_function_name", ["identity", "sqrt"]) +@pytest.mark.parametrize("centering", [False, True]) +@pytest.mark.parametrize("invariant", [False, True]) def test__scaling_equivariance( - size, is_time_series, is_voxel, activation, n_output_feature, - dimension, norm_function_name, - centering, invariant, + size, + is_time_series, + is_voxel, + activation, + n_output_feature, + dimension, + norm_function_name, + centering, + invariant, ): orthogonal_tensor = PhlowerTensor(torch.eye(3)) dict_scaling_factor = { - k: - np.random.rand() * 10 for k, v in dimension.items()} + k: np.random.rand() * 10 for k, v in dimension.items() + } scaling_factor = np.prod( - [dict_scaling_factor[k]**v for k, v in dimension.items()]) + [dict_scaling_factor[k] ** v for k, v in dimension.items()] + ) create_linear_weight = size[-1] != n_output_feature model = SimilarityEquivariantMLP( @@ -261,37 +298,53 @@ def test__scaling_equivariance( create_linear_weight=create_linear_weight, norm_function_name=norm_function_name, disable_en_equivariance=True, - centering=centering, invariant=invariant, + centering=centering, + invariant=invariant, ) t = phlower_tensor( - torch.rand(*size) * 2 - 1., dimension=dimension, - is_time_series=is_time_series, is_voxel=is_voxel) + torch.rand(*size) * 2 - 1.0, + dimension=dimension, + is_time_series=is_time_series, + is_voxel=is_voxel, + ) dict_tensor = {"tensor": t} dict_scales = { k: phlower_tensor( - torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1})) + torch.rand(1) * 10, dimension=PhysicalDimensions({k: 1}) + ) for k, v in dimension.items() } ts = phlower_tensor_collection(dict_tensor) if invariant: - actual_tensor = _functions.apply_orthogonal_group( - orthogonal_tensor, model(ts, dict_scales=dict_scales)) * 1. + actual_tensor = ( + _functions.apply_orthogonal_group( + orthogonal_tensor, model(ts, dict_scales=dict_scales) + ) + * 1.0 + ) else: - actual_tensor = _functions.apply_orthogonal_group( - orthogonal_tensor, - model(ts, dict_scales=dict_scales)) * scaling_factor + actual_tensor = ( + _functions.apply_orthogonal_group( + orthogonal_tensor, model(ts, dict_scales=dict_scales) + ) + * scaling_factor + ) actual = actual_tensor.to_numpy() - dict_transformed = {'tensor': _functions.apply_orthogonal_group( - orthogonal_tensor, t) * scaling_factor} + dict_transformed = { + "tensor": _functions.apply_orthogonal_group(orthogonal_tensor, t) + * scaling_factor + } dict_scaled_scales = { - k: v * dict_scaling_factor[k] for k, v in dict_scales.items()} + k: v * dict_scaling_factor[k] for k, v in dict_scales.items() + } transformed_phlower_tensors = phlower_tensor_collection(dict_transformed) desired = model( - transformed_phlower_tensors, dict_scales=dict_scaled_scales).to_numpy() + transformed_phlower_tensors, dict_scales=dict_scaled_scales + ).to_numpy() scale = np.max(np.abs(desired)) # Test equivariance @@ -300,7 +353,7 @@ def test__scaling_equivariance( if invariant: # Test dimensionless in case of invariant for v in actual_tensor.dimension.to_dict().values(): - np.testing.assert_almost_equal(v, 0.) + np.testing.assert_almost_equal(v, 0.0) else: # Test dimension is kept for k, v in actual_tensor.dimension.to_dict().items(): @@ -312,12 +365,14 @@ def test__similarity_equivariance_no_dimension(): t = phlower_tensor(torch.rand(10, 3, 8), dimension=None) t_scale = phlower_tensor( - torch.tensor(0.1), dimension=PhysicalDimensions({"T": 1})) + torch.tensor(0.1), dimension=PhysicalDimensions({"T": 1}) + ) dict_scales = {"T": t_scale} ts = phlower_tensor_collection({"tensor": t}) with pytest.raises(PhlowerDimensionRequiredError): model(ts, dict_scales=dict_scales) + def test__similarity_equivariance_no_scale_input(): model = SimilarityEquivariantMLP(nodes=[8, 8]) dimension = {"T": -1, "L": 1, "M": 0, "I": 0, "Theta": 0, "N": 0, "J": 0} From 5892d03809bad8eb0dc6e0b767c75d46802bf9bf Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 14 Aug 2024 18:12:59 +0900 Subject: [PATCH 45/89] add new stage to test document --- .gitlab-ci.yml | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 3038e63..5e2d5f6 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -1,5 +1,6 @@ stages: - test + - document_test - deploy default: @@ -45,15 +46,27 @@ e2e_test: - no-gpu - GenuineIntel +pages_test: + stage: document_test + script: + - make document + tags: + - no-gpu + - GenuineIntel + artifacts: + paths: + - docs/build/html/ + pages: stage: deploy script: - - make document - mkdir public - cp -r docs/build/html/* public/ artifacts: paths: - public + dependencies: + - pages_test only: - main - develop From 2769cf10f53b6d133bc2d6364fa24ff20a41c960 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 14 Aug 2024 18:18:22 +0900 Subject: [PATCH 46/89] update test --- tests/test_nn/test_core_modules/test_pinv.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_nn/test_core_modules/test_pinv.py b/tests/test_nn/test_core_modules/test_pinv.py index cf7bdde..a9cabe5 100644 --- a/tests/test_nn/test_core_modules/test_pinv.py +++ b/tests/test_nn/test_core_modules/test_pinv.py @@ -19,10 +19,10 @@ def test__can_call_parameters(): @pytest.mark.parametrize( "mlp_nodes, activations, decimal", [ - ([10, 10], ["identity"], 4), - ([10, 12], ["leaky_relu0p5"], 4), - ([20, 40, 100], ["tanh", "identity"], 4), - ([20, 20, 40, 100], ["tanh", "smooth_leaky_relu", "leaky_relu0p5"], 3), + ([10, 10], ["identity"], 3), + ([10, 12], ["leaky_relu0p5"], 3), + ([20, 40, 100], ["tanh", "identity"], 3), + ([20, 20, 40, 100], ["tanh", "smooth_leaky_relu", "leaky_relu0p5"], 2), ], ) def test__pinv_mlp(mlp_nodes, activations, decimal): From 6087db619fca807b7bcf4af2fab13cea7d1e7338 Mon Sep 17 00:00:00 2001 From: horiem Date: Fri, 16 Aug 2024 15:33:08 +0900 Subject: [PATCH 47/89] reflect reviews further --- src/phlower/nn/_core_modules/__init__.py | 11 ++++------- .../nn/_core_modules/_en_equivariant_mlp.py | 6 ++++-- .../_core_modules/_similarity_equivariant_mlp.py | 14 ++++++-------- src/phlower/nn/_core_modules/_utils.py | 4 +--- 4 files changed, 15 insertions(+), 20 deletions(-) diff --git a/src/phlower/nn/_core_modules/__init__.py b/src/phlower/nn/_core_modules/__init__.py index 78f8ed6..4fbc6d0 100644 --- a/src/phlower/nn/_core_modules/__init__.py +++ b/src/phlower/nn/_core_modules/__init__.py @@ -1,19 +1,16 @@ from phlower.nn._core_modules._concatenator import Concatenator +from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._identity import Identity from phlower.nn._core_modules._mlp import MLP from phlower.nn._core_modules._pinv_mlp import PInvMLP from phlower.nn._core_modules._proportional import Proportional from phlower.nn._core_modules._share import Share +from phlower.nn._core_modules._similarity_equivariant_mlp import ( + SimilarityEquivariantMLP, +) from phlower.nn._interface_module import IPhlowerCoreModule -if True: - # NOTE: Import advanced models after - from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP - from phlower.nn._core_modules._similarity_equivariant_mlp import ( - SimilarityEquivariantMLP, - ) - _all_models: list[IPhlowerCoreModule] = [ Concatenator, EnEquivariantMLP, diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py index 5d67649..2482aa8 100644 --- a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py @@ -5,7 +5,9 @@ from phlower._base.tensors import PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections -from phlower.nn._core_modules import Identity, Proportional, _functions, _utils +from phlower.nn._core_modules import _functions, _utils +from phlower.nn._core_modules._identity import Identity +from phlower.nn._core_modules._proportional import Proportional from phlower.nn._interface_module import ( IPhlowerCoreModule, IReadonlyReferenceGroup, @@ -48,7 +50,7 @@ def __init__( dropouts: list[float] | None = None, bias: bool = False, create_linear_weight: bool = False, - norm_function_name: str = None, + norm_function_name: str = "identity", ) -> None: super().__init__() diff --git a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py index 750c943..4ee02ae 100644 --- a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py @@ -9,13 +9,11 @@ IPhlowerTensorCollections, phlower_tensor_collection, ) -from phlower.nn._core_modules import ( - MLP, - EnEquivariantMLP, - Identity, - Proportional, - _functions, -) +from phlower.nn._core_modules import _functions +from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP +from phlower.nn._core_modules._identity import Identity +from phlower.nn._core_modules._mlp import MLP +from phlower.nn._core_modules._proportional import Proportional from phlower.nn._interface_module import ( IPhlowerCoreModule, IReadonlyReferenceGroup, @@ -66,7 +64,7 @@ def __init__( dropouts: list[float] | None = None, bias: bool = False, create_linear_weight: bool = False, - norm_function_name: str = None, + norm_function_name: str = "identity", disable_en_equivariance: bool = False, invariant: bool = False, centering: bool = False, diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index 23b8b20..e7b7855 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -99,10 +99,8 @@ def select(name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: @staticmethod def select_inverse( - name: str | None, + name: str, ) -> Callable[[torch.Tensor], torch.Tensor]: - if name is None: - name = "identity" return ActivationSelector._REGISTERED_ACTIVATIONS[ ActivationSelector._inverse_activation_name(name) ] From 7bda12499b66af526b6e41a42f54696005345189 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 19 Aug 2024 11:06:12 +0900 Subject: [PATCH 48/89] [fix] fix select function in ActivationSelector --- src/phlower/nn/_core_modules/_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index e7b7855..78382bd 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -92,9 +92,7 @@ class ActivationSelector: } @staticmethod - def select(name: str | None) -> Callable[[torch.Tensor], torch.Tensor]: - if name is None: - name = "identity" + def select(name: str) -> Callable[[torch.Tensor], torch.Tensor]: return ActivationSelector._REGISTERED_ACTIVATIONS[name] @staticmethod From 6a89e2bf1e4e51ade9921506f128aba651796642 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 19 Aug 2024 11:36:58 +0900 Subject: [PATCH 49/89] change ruff settings. add ANN linter --- Makefile | 2 +- pyproject.toml | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/Makefile b/Makefile index e1d8a55..1380d3e 100644 --- a/Makefile +++ b/Makefile @@ -14,7 +14,7 @@ mypy: .PHONY: format format: - poetry run python3 -m ruff check --select I --fix + poetry run python3 -m ruff check --fix poetry run python3 -m ruff format .PHONY: test diff --git a/pyproject.toml b/pyproject.toml index 2c3f4c7..6613865 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,6 +104,7 @@ select = [ "B", # flake8-bugbear "C4", # flake8-comprehensions "UP", # pyupgrade + "ANN", # flake8-annotations ] ignore = [] From 16f84983b56df14f3740dfb8cbf9322d697dc167 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 19 Aug 2024 14:04:08 +0900 Subject: [PATCH 50/89] add type annotations for all functions --- pyproject.toml | 8 +- src/phlower/_base/_dimension.py | 9 +-- src/phlower/_base/_functionals/_check.py | 4 +- .../_base/_functionals/_concatenate.py | 8 +- src/phlower/_base/_functionals/_to_batch.py | 2 +- .../array/sparse/_sparse_array_wrapper.py | 20 +++-- .../_base/tensors/_dimension_tensor.py | 79 ++++++++++++++----- src/phlower/_base/tensors/_interface.py | 4 +- src/phlower/_base/tensors/_phlower_tensor.py | 30 +++---- .../collections/arrays/_arrays_dict.py | 6 +- .../tensors/_tensor_collections.py | 4 +- src/phlower/io/_directory.py | 2 +- src/phlower/io/_files/checkpoint_file.py | 4 +- src/phlower/io/_files/numpy_file.py | 23 +++--- src/phlower/io/_files/yaml_file.py | 5 +- .../nn/_core_modules/_en_equivariant_mlp.py | 2 +- src/phlower/nn/_core_modules/_functions.py | 18 ++--- src/phlower/nn/_core_modules/_pinv_mlp.py | 12 +-- .../_similarity_equivariant_mlp.py | 2 +- src/phlower/nn/_group_module.py | 2 +- .../scale_functions/identity_scaler.py | 8 +- .../_scalers/scale_functions/isoam_scaler.py | 8 +- .../scale_functions/max_abs_scaler.py | 8 +- .../scale_functions/min_max_scaler.py | 12 ++- .../scale_functions/standard_scaler.py | 2 +- .../preprocessing/features/_feature_store.py | 5 +- src/phlower/services/preprocessing/scaling.py | 7 +- src/phlower/services/trainer/_optimizer.py | 2 +- .../services/trainer/_trainer_logger.py | 2 +- src/phlower/settings/_group_settings.py | 4 +- .../settings/_module_parameter_setting.py | 6 +- .../settings/_module_settings/__init__.py | 2 +- .../settings/_module_settings/_gcn_setting.py | 2 +- .../settings/_module_settings/_mlp_setting.py | 2 +- src/phlower/settings/_phlower_setting.py | 2 +- src/phlower/settings/_scaling_setting.py | 14 ++-- src/phlower/settings/_trainer_setting.py | 4 +- src/phlower/utils/_logging.py | 4 +- src/phlower/utils/_optimizer.py | 2 +- src/phlower/utils/_schedulers.py | 2 +- src/phlower/utils/preprocess.py | 2 +- tests/e2e_tests/test_preprocess.py | 9 ++- tests/e2e_tests/test_train.py | 11 +-- .../test_array/test_sparse_array_wrapper.py | 5 +- tests/test_base/test_dimensions.py | 9 +-- tests/test_base/test_functionals/conftest.py | 3 +- .../test_base/test_functionals/test_batch.py | 19 ++++- .../test_base/test_functionals/test_check.py | 16 ++-- .../test_functionals/test_unbatch.py | 33 ++++++-- .../test_tensors/test__dimensions.py | 9 +-- .../test_tensors/test__phlower_tensor.py | 34 ++++++-- tests/test_data/conftest.py | 5 +- tests/test_data/test_data_loader.py | 27 ++++--- tests/test_data/test_datasets.py | 32 ++++---- tests/test_io/test_directory.py | 15 ++-- tests/test_io/test_files/conftest.py | 2 +- .../test_files/test_checkpoint_file.py | 19 +++-- tests/test_io/test_files/test_numpy_file.py | 26 ++++-- tests/test_io/test_files/test_yaml_file.py | 17 ++-- tests/test_io/test_model_selector.py | 17 ++-- .../test_core_modules/test_concatenator.py | 5 +- .../test_en_equivariant_mlp.py | 10 ++- .../test_core_modules/test_functions.py | 63 +++++++++------ tests/test_nn/test_core_modules/test_gcn.py | 3 +- .../test_core_modules/test_identity.py | 3 +- tests/test_nn/test_core_modules/test_pinv.py | 3 +- .../test_core_modules/test_proportional.py | 5 +- tests/test_nn/test_core_modules/test_share.py | 5 +- .../test_similarity_equivariant_mlp.py | 53 ++++++------- tests/test_nn/test_group_module.py | 9 ++- .../test_identity_scale.py | 9 +-- .../test_scale_functions/test_isoam_scaler.py | 13 +-- .../test_max_abs_powered_scaler.py | 30 ++++--- .../test_min_max_scaler.py | 7 +- .../test_standard_scaler.py | 13 ++- .../test_scalers/test_scaler_wrapper.py | 12 ++- .../test_scalers/test_scalers_composition.py | 9 ++- .../test_logging_items/test_logging_items.py | 32 +++++--- .../test_trainer/test_optimizer.py | 14 +++- tests/test_settings/test_model_settings.py | 11 ++- .../test_concatenator_setting.py | 16 ++-- .../test_module_settings/test_gcn_settings.py | 26 +++--- .../test_module_settings/test_mlp_setting.py | 26 +++--- .../test_share_setting.py | 7 +- tests/test_settings/test_phlower_setting.py | 7 +- tests/test_settings/test_scaling_setting.py | 59 +++++++++----- tests/test_utils/test_encryption.py | 5 +- tests/test_utils/test_env.py | 5 +- tests/test_utils/test_multiprocessor.py | 33 +++++--- tests/test_utils/test_optimizer.py | 7 +- tests/test_utils/test_preprocess.py | 11 ++- tests/test_utils/test_progress_bar.py | 7 +- tests/test_utils/test_schedulers.py | 7 +- tests/test_utils/test_timer.py | 11 ++- 94 files changed, 714 insertions(+), 474 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6613865..6243bc6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.ruff] -src = ["src", "tests"] +include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"] # Exclude a variety of commonly ignored directories. exclude = [ @@ -106,7 +106,7 @@ select = [ "UP", # pyupgrade "ANN", # flake8-annotations ] -ignore = [] +ignore = ["ANN003", "ANN101", "ANN102", "ANN204"] # Allow fix for all enabled rules (when `--fix`) is provided. fixable = ["ALL"] @@ -115,8 +115,12 @@ unfixable = [] # Allow unused variables when underscore-prefixed. dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" +[tool.ruff.lint.flake8-annotations] +suppress-none-returning = true + [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F401"] +"src/phlower/_base/**/*.py" = ["ANN401"] [tool.ruff.format] # Like Black, use double quotes for strings. diff --git a/src/phlower/_base/_dimension.py b/src/phlower/_base/_dimension.py index 1834025..fc089dd 100644 --- a/src/phlower/_base/_dimension.py +++ b/src/phlower/_base/_dimension.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Annotated, Any +from typing import Annotated from pydantic import ( PlainSerializer, @@ -13,13 +13,12 @@ from phlower.utils.exceptions import InvalidDimensionError -def _validate(v: Any, info: ValidationInfo): +def _validate(v: dict, info: ValidationInfo) -> PhysicalDimensions: if not isinstance(v, dict): raise TypeError(f"Expected dictionary, but got {type(v)}") try: - ans = PhysicalDimensions(v) - return ans + return PhysicalDimensions(v) except Exception as ex: raise TypeError("Validation for physical dimension is failed.") from ex @@ -78,7 +77,7 @@ def to_list(self) -> list[float]: return _list - def to_dict(self): + def to_dict(self) -> dict: return self._dimensions diff --git a/src/phlower/_base/_functionals/_check.py b/src/phlower/_base/_functionals/_check.py index 46d48bc..6412a0d 100644 --- a/src/phlower/_base/_functionals/_check.py +++ b/src/phlower/_base/_functionals/_check.py @@ -6,13 +6,13 @@ from phlower._base.tensors._interface import IPhlowerTensor -def is_same_layout(tensors: Sequence[IPhlowerTensor]): +def is_same_layout(tensors: Sequence[IPhlowerTensor]) -> bool: is_sparse_flags: set[bool] = set(tensors | select(lambda x: x.is_sparse)) return len(is_sparse_flags) == 1 -def is_same_dimensions(tensors: Sequence[IPhlowerTensor]): +def is_same_dimensions(tensors: Sequence[IPhlowerTensor]) -> bool: if len(tensors) == 0: return True diff --git a/src/phlower/_base/_functionals/_concatenate.py b/src/phlower/_base/_functionals/_concatenate.py index 2865e35..35c578f 100644 --- a/src/phlower/_base/_functionals/_concatenate.py +++ b/src/phlower/_base/_functionals/_concatenate.py @@ -9,7 +9,9 @@ from ._check import is_same_dimensions, is_same_layout -def concatenate(tensors: Sequence[IPhlowerTensor], dense_dim: int = 0): +def concatenate( + tensors: Sequence[IPhlowerTensor], dense_dim: int = 0 +) -> IPhlowerTensor: if not is_same_layout(tensors): raise ValueError("Cannot concatenate dense tensor and sparse tensor") @@ -25,7 +27,9 @@ def concatenate(tensors: Sequence[IPhlowerTensor], dense_dim: int = 0): return _dense_concatenate(tensors, dim=dense_dim) -def _dense_concatenate(tensors: Sequence[IPhlowerTensor], dim: int = 0): +def _dense_concatenate( + tensors: Sequence[IPhlowerTensor], dim: int = 0 +) -> IPhlowerTensor: return torch.concatenate(tensors, dim=dim) diff --git a/src/phlower/_base/_functionals/_to_batch.py b/src/phlower/_base/_functionals/_to_batch.py index 4ef8403..1c7f74f 100644 --- a/src/phlower/_base/_functionals/_to_batch.py +++ b/src/phlower/_base/_functionals/_to_batch.py @@ -26,7 +26,7 @@ def to_batch( def _create_batch_info( tensors: Sequence[IPhlowerTensor], dense_concat_dim: int -): +) -> GraphBatchInfo: _shapes = list(tensors | select(lambda x: x.shape)) if tensors[0].is_sparse: diff --git a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py index adef209..ffbb88c 100644 --- a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py +++ b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py @@ -25,19 +25,19 @@ def is_sparse(self) -> bool: return True @property - def is_time_series(self): + def is_time_series(self) -> bool: return False @property - def row(self): + def row(self) -> np.ndarray: return self._sparse_data.row @property - def col(self): + def col(self) -> np.ndarray: return self._sparse_data.col @property - def data(self): + def data(self) -> np.ndarray: return self._sparse_data.data @property @@ -138,12 +138,16 @@ def batch( return SparseArrayWrapper(concat_arr), info -def unbatch(array: SparseArrayWrapper, batch_info: GraphBatchInfo): +def unbatch( + array: SparseArrayWrapper, batch_info: GraphBatchInfo +) -> list[SparseArrayWrapper]: results = _sparse_decompose(array.to_numpy(), batch_info) return [SparseArrayWrapper(arr) for arr in results] -def _sparse_concatenate(arrays: Sequence[SparseArrayWrapper]): +def _sparse_concatenate( + arrays: Sequence[SparseArrayWrapper], +) -> SparseArrayWrapper: offsets = np.cumsum( np.array( [ @@ -170,7 +174,9 @@ def _sparse_concatenate(arrays: Sequence[SparseArrayWrapper]): return sparse_arr -def _sparse_decompose(array: SparseArrayType, batch_info: GraphBatchInfo): +def _sparse_decompose( + array: SparseArrayType, batch_info: GraphBatchInfo +) -> list[sp.coo_matrix]: sizes = np.cumsum(batch_info.sizes) offsets = np.cumsum( np.array([[0, 0]] + batch_info.shapes[:-1], dtype=np.int32), axis=0 diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index e8bdfd0..af3bbc5 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -2,6 +2,7 @@ import functools from collections.abc import Callable +from typing import Any import torch @@ -97,7 +98,7 @@ def __repr__(self) -> str: ) return f"{self.__class__.__name__}({texts})" - def to_physics_dimension(self): + def to_physics_dimension(self) -> PhysicalDimensions: _dict = self.to_dict() return PhysicalDimensions(_dict) @@ -118,7 +119,13 @@ def is_dimensionless(self) -> bool: return torch.sum(torch.abs(self._tensor)) < 1e-5 @classmethod - def __torch_function__(cls, func, types, args: tuple, kwargs=None): + def __torch_function__( + cls, + func: Callable, + types: list[type], + args: tuple, + kwargs: dict | None = None, + ) -> PhlowerDimensionTensor: if kwargs is None: kwargs = {} @@ -128,10 +135,10 @@ def __torch_function__(cls, func, types, args: tuple, kwargs=None): return _HANDLED_FUNCTIONS[func](*args, **kwargs) -def dimension_wrap_implements(torch_function): +def dimension_wrap_implements(torch_function: Callable) -> Callable: """Register a torch function override for PhysicsUnitTensor""" - def decorator(func): + def decorator(func: Callable) -> Callable: functools.update_wrapper(func, torch_function) _HANDLED_FUNCTIONS[torch_function] = func return func @@ -140,12 +147,14 @@ def decorator(func): @dimension_wrap_implements(torch.mean) -def mean(inputs: PhlowerDimensionTensor): +def mean(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(inputs._tensor) @dimension_wrap_implements(torch.add) -def add(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): +def add( + inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor +) -> PhlowerDimensionTensor: if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): if inputs != other: raise DimensionIncompatibleError( @@ -159,7 +168,9 @@ def add(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): @dimension_wrap_implements(torch.sub) -def sub(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): +def sub( + inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor +) -> PhlowerDimensionTensor: if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): if inputs != other: raise DimensionIncompatibleError( @@ -173,17 +184,22 @@ def sub(inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor): @dimension_wrap_implements(torch.pow) -def pow(inputs: PhlowerDimensionTensor, other): +def pow( + inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor +) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(inputs._tensor * other) @dimension_wrap_implements(torch.sqrt) -def sqrt(inputs: PhlowerDimensionTensor): +def sqrt(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(inputs._tensor / 2) @dimension_wrap_implements(torch.mul) -def mul(inputs, other): +def mul( + inputs: PhlowerDimensionTensor | torch.Tensor, + other: PhlowerDimensionTensor | torch.Tensor, +) -> PhlowerDimensionTensor: _input = ( inputs if isinstance(inputs, PhlowerDimensionTensor) @@ -198,7 +214,9 @@ def mul(inputs, other): @dimension_wrap_implements(torch.div) -def div(inputs, other): +def div( + inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor +) -> PhlowerDimensionTensor: _input = ( inputs if isinstance(inputs, PhlowerDimensionTensor) @@ -213,7 +231,9 @@ def div(inputs, other): @dimension_wrap_implements(torch.reshape) -def reshape(inputs, shape): +def reshape( + inputs: PhlowerDimensionTensor, shape: tuple[int] +) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(inputs._tensor) @@ -233,7 +253,9 @@ def cat( @dimension_wrap_implements(torch.sparse.mm) -def sparse_mm(inputs, other): +def sparse_mm( + inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor +) -> PhlowerDimensionTensor: if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): return PhlowerDimensionTensor(inputs._tensor + other._tensor) @@ -241,17 +263,21 @@ def sparse_mm(inputs, other): @dimension_wrap_implements(torch.nn.functional.linear) -def nn_linear(inputs, *args): +def nn_linear( + inputs: PhlowerDimensionTensor, *args: Any +) -> PhlowerDimensionTensor: return inputs @dimension_wrap_implements(torch.nn.functional.dropout) -def dropout(inputs, *args, **kwards): +def dropout( + inputs: PhlowerDimensionTensor, *args: Any, **kwards: Any +) -> PhlowerDimensionTensor: return inputs @dimension_wrap_implements(torch.stack) -def stack(inputs): +def stack(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: if all(isinstance(v, PhlowerDimensionTensor) for v in inputs): # HACK: is it possible to use unique method ? for v in inputs: @@ -264,7 +290,12 @@ def stack(inputs): @dimension_wrap_implements(torch.nn.functional.mse_loss) -def mse_loss(inputs, others, *args, **kwards): +def mse_loss( + inputs: PhlowerDimensionTensor, + others: PhlowerDimensionTensor, + *args: Any, + **kwards: Any, +) -> PhlowerDimensionTensor: if inputs != others: raise DimensionIncompatibleError() @@ -272,7 +303,9 @@ def mse_loss(inputs, others, *args, **kwards): @dimension_wrap_implements(torch.sum) -def _sum(inputs, *args, **kwards): +def _sum( + inputs: PhlowerDimensionTensor, *args: Any, **kwards: Any +) -> PhlowerDimensionTensor: if isinstance(inputs, PhlowerDimensionTensor): return inputs @@ -280,7 +313,9 @@ def _sum(inputs, *args, **kwards): @dimension_wrap_implements(torch.concatenate) -def concatenate(inputs, *args, **kwards): +def concatenate( + inputs: PhlowerDimensionTensor, *args: Any, **kwards: Any +) -> PhlowerDimensionTensor: if all(isinstance(v, PhlowerDimensionTensor) for v in inputs): # HACK: is it possible to use unique method ? for v in inputs: @@ -293,7 +328,7 @@ def concatenate(inputs, *args, **kwards): @dimension_wrap_implements(torch.tanh) -def tanh(tensor: PhlowerDimensionTensor): +def tanh(tensor: PhlowerDimensionTensor) -> PhlowerDimensionTensor: if not tensor.is_dimensionless: raise DimensionIncompatibleError( f"Should be dimensionless to apply tanh but {tensor}" @@ -302,7 +337,9 @@ def tanh(tensor: PhlowerDimensionTensor): @dimension_wrap_implements(torch.nn.functional.leaky_relu) -def leaky_relu(tensor: PhlowerDimensionTensor, *args, **kwargs): +def leaky_relu( + tensor: PhlowerDimensionTensor, *args: Any, **kwargs: Any +) -> PhlowerDimensionTensor: if not tensor.is_dimensionless: raise DimensionIncompatibleError( f"Should be dimensionless to apply leaky_relu but {tensor}" diff --git a/src/phlower/_base/tensors/_interface.py b/src/phlower/_base/tensors/_interface.py index 8476273..a9a429c 100644 --- a/src/phlower/_base/tensors/_interface.py +++ b/src/phlower/_base/tensors/_interface.py @@ -40,10 +40,10 @@ def values(self) -> torch.Tensor: ... def __str__(self) -> str: ... @abc.abstractmethod - def __add__(self, other) -> IPhlowerTensor: ... + def __add__(self, other: IPhlowerTensor) -> IPhlowerTensor: ... @abc.abstractmethod - def __mul__(self, other) -> IPhlowerTensor: ... + def __mul__(self, other: IPhlowerTensor) -> IPhlowerTensor: ... @abc.abstractmethod def __len__(self) -> int: ... diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index ac68b33..0b5655a 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -41,7 +41,7 @@ def phlower_tensor( ) = None, is_time_series: bool = False, is_voxel: bool = False, -): +) -> PhlowerTensor: if isinstance(tensor, PhlowerTensor): if dimension is not None: logger.warning( @@ -143,19 +143,19 @@ def __str__(self) -> str: f"Dimension: {self._dimension_tensor})" ) - def __eq__(self, other): + def __eq__(self, other: PhlowerTensor): return torch.eq(self, other) - def __lt__(self, other): + def __lt__(self, other: PhlowerTensor): return torch.lt(self, other) - def __le__(self, other): + def __le__(self, other: PhlowerTensor): return torch.le(self, other) - def __gt__(self, other): + def __gt__(self, other: PhlowerTensor): return torch.gt(self, other) - def __ge__(self, other): + def __ge__(self, other: PhlowerTensor): return torch.ge(self, other) def __abs__(self) -> PhlowerTensor: @@ -167,28 +167,28 @@ def __sub__(self, other: PhlowerTensor): def __neg__(self): return torch.neg(self) - def __add__(self, other) -> PhlowerTensor: + def __add__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.add(self, other) - def __radd__(self, other) -> PhlowerTensor: + def __radd__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.add(self, other) - def __mul__(self, other) -> PhlowerTensor: + def __mul__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.mul(self, other) - def __rmul__(self, other) -> PhlowerTensor: + def __rmul__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.mul(self, other) - def __truediv__(self, other) -> PhlowerTensor: + def __truediv__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.div(self, other) - def __rtruediv__(self, other) -> PhlowerTensor: + def __rtruediv__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.div(other, self) - def __pow__(self, other) -> PhlowerTensor: + def __pow__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.pow(self, other) - def __setitem__(self, key, value): + def __setitem__(self, key: str, value: float): if isinstance(key, PhlowerTensor): self._tensor[key.to_tensor()] = value else: @@ -365,7 +365,7 @@ def __torch_function__( types: list[type], args: tuple, kwargs: dict | None = None, - ): + ) -> PhlowerTensor: if func.__name__ in UNSUPPORTED_FUNCTION_NAMES: raise PhlowerUnsupportedTorchFunctionError( f"Unsupported function: {func.__name__}" diff --git a/src/phlower/collections/arrays/_arrays_dict.py b/src/phlower/collections/arrays/_arrays_dict.py index 8088114..df6cea6 100644 --- a/src/phlower/collections/arrays/_arrays_dict.py +++ b/src/phlower/collections/arrays/_arrays_dict.py @@ -22,7 +22,7 @@ def __init__(self, name: str, data: list[IPhlowerArray]) -> None: self._is_sparse = self._reduce_is_sparse() self._is_time_series = self._reduce_is_time_series() - def _reduce_is_sparse(self): + def _reduce_is_sparse(self) -> bool: _is_sparse = np.unique(np.array([v.is_sparse for v in self._data])) if len(_is_sparse) != 1: raise ValueError( @@ -31,7 +31,7 @@ def _reduce_is_sparse(self): ) return _is_sparse.item() - def _reduce_is_time_series(self): + def _reduce_is_time_series(self) -> bool: _is_time_series = np.unique( np.array([v.is_time_series for v in self._data]) ) @@ -60,7 +60,7 @@ def to_batched_tensor( dimensions: PhysicalDimensions | None = None, is_time_series: bool = False, is_voxel: bool = False, - ): + ) -> PhlowerTensor: tensors = [ v.to_phlower_tensor( device=device, diff --git a/src/phlower/collections/tensors/_tensor_collections.py b/src/phlower/collections/tensors/_tensor_collections.py index 473ac76..113e223 100644 --- a/src/phlower/collections/tensors/_tensor_collections.py +++ b/src/phlower/collections/tensors/_tensor_collections.py @@ -89,7 +89,7 @@ def __contains__(self, key: str): def __len__(self) -> int: return len(self._x) - def values(self): + def values(self) -> Iterable[dict[str, Any]]: return self._x.values() def keys(self) -> Iterable[str]: @@ -101,7 +101,7 @@ def items(self) -> abc.ItemsView: def pop(self, key: str, default: PhlowerTensor = None) -> PhlowerTensor: return self._x.pop(key, default) - def __getitem__(self, key: Any) -> PhlowerTensor: + def __getitem__(self, key: str) -> PhlowerTensor: if isinstance(key, str): return self._x[key] diff --git a/src/phlower/io/_directory.py b/src/phlower/io/_directory.py index e95e379..edd625e 100644 --- a/src/phlower/io/_directory.py +++ b/src/phlower/io/_directory.py @@ -91,7 +91,7 @@ def _find_file( builder: Callable[[pathlib.Path], Any] | None = None, *, allow_missing: bool = False, - ): + ) -> pathlib.Path: if builder is None: builder = lambda x: x # NOQA diff --git a/src/phlower/io/_files/checkpoint_file.py b/src/phlower/io/_files/checkpoint_file.py index d18e69e..18e47d2 100644 --- a/src/phlower/io/_files/checkpoint_file.py +++ b/src/phlower/io/_files/checkpoint_file.py @@ -89,13 +89,13 @@ def epoch(self) -> int: def is_encrypted(self) -> bool: return self._ext_type == PhlowerFileExtType.PTHENC - def load(self, device: str, *, decrypt_key: bytes = None) -> Any: + def load(self, device: str, *, decrypt_key: bytes = None) -> Any: # noqa: ANN401 if self.is_encrypted: return self._load_encrypted(device=device, decrypt_key=decrypt_key) else: return self._load(device=device) - def _load(self, device: str) -> Any: + def _load(self, device: str) -> Any: # noqa: ANN401 return torch.load(self._path, map_location=device) def _load_encrypted( diff --git a/src/phlower/io/_files/numpy_file.py b/src/phlower/io/_files/numpy_file.py index 1f33f21..f646bd4 100644 --- a/src/phlower/io/_files/numpy_file.py +++ b/src/phlower/io/_files/numpy_file.py @@ -4,7 +4,7 @@ import io import os import pathlib -from typing import get_args +from typing import Any, get_args import numpy as np import scipy.sparse as sp @@ -95,7 +95,7 @@ def is_encrypted(self) -> bool: return False @property - def file_extension(self): + def file_extension(self) -> str: return self._ext_type.value @property @@ -117,7 +117,9 @@ def get_variable_name(self) -> str: return name -def _get_fileio(data: ArrayDataType, encrypt_key: bytes | None = None): +def _get_fileio( + data: ArrayDataType, encrypt_key: bytes | None = None +) -> INumpyFileIOCore: if isinstance(data, np.ndarray): if encrypt_key is not None: return _NpyEncFileIO() @@ -134,15 +136,18 @@ def _get_fileio(data: ArrayDataType, encrypt_key: bytes | None = None): class INumpyFileIOCore(metaclass=abc.ABCMeta): - @abc.abstractclassmethod + @classmethod + @abc.abstractmethod def load(cls, path: pathlib.Path, decrypt_key: bytes = None): ... - @abc.abstractclassmethod + @classmethod + @abc.abstractmethod def get_save_path( cls, output_directory: pathlib.Path, file_basename: str ) -> pathlib.Path: ... - @abc.abstractclassmethod + @classmethod + @abc.abstractmethod def save( cls, save_path: pathlib.Path, @@ -179,7 +184,7 @@ def save( class _NpyEncFileIO(INumpyFileIOCore): @classmethod - def load(cls, path: pathlib.Path, decrypt_key: bytes = None): + def load(cls, path: pathlib.Path, decrypt_key: bytes = None) -> Any: # noqa: ANN401 if decrypt_key is None: raise ValueError( "encrpt key is None. Cannot decrypt encrypted file." @@ -216,7 +221,7 @@ def save( class _NpzFileIO(INumpyFileIOCore): @classmethod - def load(cls, path: pathlib.Path, decrypt_key: bytes = None): + def load(cls, path: pathlib.Path, decrypt_key: bytes = None) -> Any: # noqa: ANN401 return sp.load_npz(path) @classmethod @@ -239,7 +244,7 @@ def save( class _NpzEncFileIO(INumpyFileIOCore): @classmethod - def load(cls, path: pathlib.Path, decrypt_key: bytes = None): + def load(cls, path: pathlib.Path, decrypt_key: bytes = None) -> Any: # noqa: ANN401: if decrypt_key is None: raise ValueError("Key is None. Cannot decrypt encrypted file.") diff --git a/src/phlower/io/_files/yaml_file.py b/src/phlower/io/_files/yaml_file.py index 74ff08b..83d6c36 100644 --- a/src/phlower/io/_files/yaml_file.py +++ b/src/phlower/io/_files/yaml_file.py @@ -1,6 +1,5 @@ import io import pathlib -from typing import Any import yaml from typing_extensions import Self @@ -113,7 +112,7 @@ def _save(file_path: pathlib.Path, dump_data: object) -> None: yaml.dump(dump_data, fw) -def _load_yaml_file(file_path: str | pathlib.Path): +def _load_yaml_file(file_path: str | pathlib.Path) -> dict: """Load YAML file. Parameters @@ -131,7 +130,7 @@ def _load_yaml_file(file_path: str | pathlib.Path): return dict_data -def _load_yaml(source: str | pathlib.Path | io.TextIOBase | Any): +def _load_yaml(source: str | pathlib.Path | io.TextIOBase) -> dict: """Load YAML source. Parameters diff --git a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py index 2482aa8..e817c6b 100644 --- a/src/phlower/nn/_core_modules/_en_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_en_equivariant_mlp.py @@ -72,7 +72,7 @@ def __init__( self._norm_function_name ) - def _init_linear_weight(self): + def _init_linear_weight(self) -> Identity | Proportional: if not self._create_linear_weight: if self._nodes[0] != self._nodes[-1]: raise ValueError( diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 09eab0e..5f14f6f 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -7,21 +7,21 @@ from phlower.utils.exceptions import PhlowerIncompatibleTensorError -def identity(x): +def identity(x: torch.Tensor) -> torch.Tensor: return x -def leaky_relu0p5(x): +def leaky_relu0p5(x: torch.Tensor) -> torch.Tensor: """Leaky ReLU with the negative slope = 0.5.""" return torch.nn.functional.leaky_relu(x, negative_slope=0.5) -def inversed_leaky_relu0p5(x): +def inversed_leaky_relu0p5(x: torch.Tensor) -> torch.Tensor: """Inverse of leaky_relu0p5.""" return torch.nn.functional.leaky_relu(x, negative_slope=2) -def truncated_atanh(x, epsilon=1e-8): +def truncated_atanh(x: torch.Tensor, epsilon: float = 1e-8) -> torch.Tensor: """Inverse tanh with truncating values >=1 or <=-1.""" x[x >= 1.0 - epsilon] = 1.0 - epsilon x[x <= -1.0 + epsilon] = -1.0 + epsilon @@ -35,11 +35,11 @@ def __init__(self, a: float = 0.75, b: float = 1 / 100): self.a = a self.b = b - def __call__(self, x): + def __call__(self, x: torch.Tensor): # a x + (1 - a) sqrt(x**2 + b) return self.a * x + (1 - self.a) * torch.sqrt(x**2 + self.b) - def inverse(self, x): + def inverse(self, x: torch.Tensor) -> torch.Tensor: return ( self.a * x - torch.sqrt( @@ -47,7 +47,7 @@ def inverse(self, x): ) ) / (2 * self.a - 1) - def derivative(self, x): + def derivative(self, x: torch.Tensor) -> torch.Tensor: return self.a - ((self.a - 1) * x) / torch.sqrt(self.b + x**2) @@ -81,7 +81,7 @@ def spmm( def einsum( - equation, + equation: str, *args: list[IPhlowerTensor], dimension: ( PhysicalDimensions @@ -251,7 +251,7 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: def tensor_times_scalar( tensor: IPhlowerTensor, scalar: int | float | IPhlowerTensor -): +) -> IPhlowerTensor: """ Compute multiplication between tensor and scalar (field). diff --git a/src/phlower/nn/_core_modules/_pinv_mlp.py b/src/phlower/nn/_core_modules/_pinv_mlp.py index 8b4e950..0017605 100644 --- a/src/phlower/nn/_core_modules/_pinv_mlp.py +++ b/src/phlower/nn/_core_modules/_pinv_mlp.py @@ -63,14 +63,16 @@ def _initialize(self): ] self._chains = self._init_pinv_chains() - def _init_pinv_chains(self): + def _init_pinv_chains(self) -> list[PInvLinear]: name = self._reference.__class__.__name__ if name in ["MLP", "Proportional"]: return self._init_pinv_mlp_chains(self._reference._chains) raise ValueError(f"Unsupported reference class: {name}") - def _init_pinv_mlp_chains(self, chains: _utils.ExtendedLinearList): + def _init_pinv_mlp_chains( + self, chains: _utils.ExtendedLinearList + ) -> list[PInvLinear]: return [PInvLinear(c) for c in chains._linears[::-1]] def forward( @@ -114,18 +116,18 @@ def __init__(self, ref_linear: torch.nn.Linear): self.ref_linear = ref_linear return - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: h = torch.nn.functional.linear(x + self.bias, self.weight) return h @property - def weight(self): + def weight(self) -> torch.Tensor: """Return pseudo inversed weight.""" w = self.ref_linear.weight return torch.pinverse(w) @property - def bias(self): + def bias(self) -> torch.Tensor | float: """Return inverse bias.""" if self.ref_linear.bias is None: return 0 diff --git a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py index 4ee02ae..1be1832 100644 --- a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py @@ -96,7 +96,7 @@ def __init__( ) self._linear_weight = self._mlp._linear_weight - def _init_linear_weight(self): + def _init_linear_weight(self) -> Identity | Proportional: if not self._create_linear_weight: if self._nodes[0] != self._nodes[-1]: raise ValueError( diff --git a/src/phlower/nn/_group_module.py b/src/phlower/nn/_group_module.py index 4f74aad..1c8d2d0 100644 --- a/src/phlower/nn/_group_module.py +++ b/src/phlower/nn/_group_module.py @@ -85,7 +85,7 @@ def get_display_info(self) -> str: def get_n_nodes(self) -> list[int]: return None - def resolve(self): + def resolve(self) -> dagstream.DagStream: for module in self._phlower_modules: module.resolve(parent=self) diff --git a/src/phlower/services/preprocessing/_scalers/scale_functions/identity_scaler.py b/src/phlower/services/preprocessing/_scalers/scale_functions/identity_scaler.py index a4ac89b..183a3b6 100644 --- a/src/phlower/services/preprocessing/_scalers/scale_functions/identity_scaler.py +++ b/src/phlower/services/preprocessing/_scalers/scale_functions/identity_scaler.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from sklearn.base import BaseEstimator, TransformerMixin from phlower.services.preprocessing._scalers import IPhlowerScaler @@ -9,7 +11,7 @@ class IdentityScaler(BaseEstimator, TransformerMixin, IPhlowerScaler): """Class to perform identity conversion (do nothing).""" @classmethod - def create(cls, name: str, **kwards): + def create(cls, name: str, **kwards) -> IdentityScaler: if name == PhlowerScalerName.IDENTITY.value: return IdentityScaler(**kwards) @@ -31,10 +33,10 @@ def is_erroneous(self) -> bool: def partial_fit(self, data: ArrayDataType) -> None: return - def transform(self, data: ArrayDataType): + def transform(self, data: ArrayDataType) -> ArrayDataType: return data - def inverse_transform(self, data: ArrayDataType): + def inverse_transform(self, data: ArrayDataType) -> ArrayDataType: return data def get_dumped_data(self) -> dict[str, str | int | float]: diff --git a/src/phlower/services/preprocessing/_scalers/scale_functions/isoam_scaler.py b/src/phlower/services/preprocessing/_scalers/scale_functions/isoam_scaler.py index a3b01fd..7f1da2b 100644 --- a/src/phlower/services/preprocessing/_scalers/scale_functions/isoam_scaler.py +++ b/src/phlower/services/preprocessing/_scalers/scale_functions/isoam_scaler.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Any import numpy as np @@ -15,7 +17,7 @@ class IsoAMScaler(BaseEstimator, TransformerMixin, IPhlowerScaler): """ @classmethod - def create(cls, name: str, **kwards): + def create(cls, name: str, **kwards) -> IsoAMScaler: if name == PhlowerScalerName.ISOAM_SCALE.value: return IsoAMScaler(**kwards) @@ -69,13 +71,13 @@ def partial_fit(self, data: ArrayDataType) -> None: self.std_ = np.sqrt(self.var_) return - def transform(self, data: ArrayDataType): + def transform(self, data: ArrayDataType) -> ArrayDataType: if self.std_ == 0.0: raise ValueError("std value is 0.") return data * (1.0 / self.std_) - def inverse_transform(self, data: ArrayDataType): + def inverse_transform(self, data: ArrayDataType) -> ArrayDataType: return data * self.std_ def get_dumped_data(self) -> dict[str, Any]: diff --git a/src/phlower/services/preprocessing/_scalers/scale_functions/max_abs_scaler.py b/src/phlower/services/preprocessing/_scalers/scale_functions/max_abs_scaler.py index cd615ac..cac5603 100644 --- a/src/phlower/services/preprocessing/_scalers/scale_functions/max_abs_scaler.py +++ b/src/phlower/services/preprocessing/_scalers/scale_functions/max_abs_scaler.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Any - import numpy as np import scipy.sparse as sp from numpy.typing import NDArray @@ -29,7 +27,7 @@ def __init__(self, power: float = 1.0, **kwargs): for k, v in kwargs.items(): setattr(self, k, self._convert(k, v)) - def _convert(self, field_name: str, value: Any): + def _convert(self, field_name: str, value: float | None) -> float: if field_name == "max_": if value is None: return value @@ -61,7 +59,7 @@ def partial_fit(self, data: ArrayDataType) -> None: def is_fitted(self) -> bool: return self.max_ is not None - def transform(self, data: ArrayDataType): + def transform(self, data: ArrayDataType) -> ArrayDataType: if not self.is_fitted(): raise ValueError( f"This scaler has not fitted yet. {self.__class__.__name__}" @@ -79,7 +77,7 @@ def transform(self, data: ArrayDataType): return data * ((1.0 / self.max_) ** self.power) - def inverse_transform(self, data: ArrayDataType): + def inverse_transform(self, data: ArrayDataType) -> ArrayDataType: if not self.is_fitted(): raise ValueError( f"This scaler has not fitted yet. {self.__class__.__name__}" diff --git a/src/phlower/services/preprocessing/_scalers/scale_functions/min_max_scaler.py b/src/phlower/services/preprocessing/_scalers/scale_functions/min_max_scaler.py index 08a3623..d768aad 100644 --- a/src/phlower/services/preprocessing/_scalers/scale_functions/min_max_scaler.py +++ b/src/phlower/services/preprocessing/_scalers/scale_functions/min_max_scaler.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Any import numpy as np @@ -10,7 +12,7 @@ class MinMaxScaler(preprocessing.MinMaxScaler, IPhlowerScaler): @classmethod - def create(cls, name: str, **kwards): + def create(cls, name: str, **kwards) -> MinMaxScaler: if name == PhlowerScalerName.MIN_MAX.value: return MinMaxScaler(**kwards) @@ -22,15 +24,17 @@ def __init__( self, feature_range: tuple[int, int] = (0, 1), *, - copy=True, - clip=False, + copy: bool = True, + clip: bool = False, **kwargs, ): super().__init__(feature_range, copy=copy, clip=clip) for k, v in kwargs.items(): setattr(self, k, self._convert(k, v)) - def _convert(self, field_name: str, value: Any): + def _convert( + self, field_name: str, value: float | str | None + ) -> np.generic | float | int | str: if field_name in "feature_range": return tuple(value) diff --git a/src/phlower/services/preprocessing/_scalers/scale_functions/standard_scaler.py b/src/phlower/services/preprocessing/_scalers/scale_functions/standard_scaler.py index 9bc523d..284bcd4 100644 --- a/src/phlower/services/preprocessing/_scalers/scale_functions/standard_scaler.py +++ b/src/phlower/services/preprocessing/_scalers/scale_functions/standard_scaler.py @@ -42,7 +42,7 @@ def __init__( for k, v in kwargs.items(): setattr(self, k, self._convert(k, v)) - def _convert(self, field_name: str, value: Any) -> Any: + def _convert(self, field_name: str, value: Any) -> Any: # noqa: ANN401 if field_name in ["copy", "with_mean", "with_std"]: return value diff --git a/src/phlower/services/preprocessing/features/_feature_store.py b/src/phlower/services/preprocessing/features/_feature_store.py index dad4307..8987085 100644 --- a/src/phlower/services/preprocessing/features/_feature_store.py +++ b/src/phlower/services/preprocessing/features/_feature_store.py @@ -1,6 +1,7 @@ from __future__ import annotations import pathlib +from collections.abc import Iterable from typing import get_args import pydantic.dataclasses as dc @@ -47,10 +48,10 @@ def _check_duplicate(self, name: str) -> None: ) return - def keys(self): + def keys(self) -> Iterable[str]: return self._feature_table.keys() - def values(self): + def values(self) -> Iterable[ArrayDataType]: return self._feature_table.values() def save( diff --git a/src/phlower/services/preprocessing/scaling.py b/src/phlower/services/preprocessing/scaling.py index 93d02e0..b19ff3c 100644 --- a/src/phlower/services/preprocessing/scaling.py +++ b/src/phlower/services/preprocessing/scaling.py @@ -98,9 +98,10 @@ def lazy_fit_all( ) # NOTE: When using multiprocessing, parameters of each scaler are - # updated in each process. However, it is not reflected in parent process. + # updated in each process. + # However, it is not reflected in parent process. # because we do not share memory. - # Thus, in this implementation, recreate scalers compositions + # Thus, in this implementation, recreate scalers compositions self._scalers.force_update(dict(results)) return @@ -127,7 +128,7 @@ def transform_file( variable_name: str, file_path: pathlib.Path | IPhlowerNumpyFile, decrypt_key: bytes | None = None, - ): + ) -> ArrayDataType: scaler_name = self._scaling_setting.get_scaler_name(variable_name) return self._scalers.transform_file( scaler_name, file_path, decrypt_key=decrypt_key diff --git a/src/phlower/services/trainer/_optimizer.py b/src/phlower/services/trainer/_optimizer.py index 61e8192..1139726 100644 --- a/src/phlower/services/trainer/_optimizer.py +++ b/src/phlower/services/trainer/_optimizer.py @@ -52,5 +52,5 @@ def step_scheduler(self): for scheduler in self._schedulers: scheduler.step() - def state_dict(self): + def state_dict(self) -> dict: return self._optimizer.state_dict() diff --git a/src/phlower/services/trainer/_trainer_logger.py b/src/phlower/services/trainer/_trainer_logger.py index 9ba8691..00668a8 100644 --- a/src/phlower/services/trainer/_trainer_logger.py +++ b/src/phlower/services/trainer/_trainer_logger.py @@ -47,7 +47,7 @@ def __init__( file_path: pathlib.Path, loss_keys: list[str] = None, display_margin: int = 4, - ): + ) -> None: self._file_path = file_path self._display_margin = display_margin self._loss_keys = loss_keys if loss_keys is not None else [] diff --git a/src/phlower/settings/_group_settings.py b/src/phlower/settings/_group_settings.py index f4228ff..380e3a5 100644 --- a/src/phlower/settings/_group_settings.py +++ b/src/phlower/settings/_group_settings.py @@ -75,7 +75,9 @@ class GroupModuleSetting( @pydantic.field_validator("modules", mode="before") @classmethod - def validate_annotate_modules(cls, vals): + def validate_annotate_modules( + cls, vals: list + ) -> list[ModuleSetting | GroupModuleSetting]: if not isinstance(vals, list): raise ValueError(f"'modules' expected to be List. actual: {vals}") diff --git a/src/phlower/settings/_module_parameter_setting.py b/src/phlower/settings/_module_parameter_setting.py index 4a04a79..afbbd4e 100644 --- a/src/phlower/settings/_module_parameter_setting.py +++ b/src/phlower/settings/_module_parameter_setting.py @@ -1,4 +1,4 @@ -from typing import Annotated +from typing import Annotated, Any import pydantic from pydantic import ( @@ -20,7 +20,9 @@ def _validate(vals: dict, info: ValidationInfo) -> IPhlowerLayerParameters: return setting_cls(**vals) -def _serialize(v: pydantic.BaseModel, info: SerializationInfo): +def _serialize( + v: pydantic.BaseModel, info: SerializationInfo +) -> dict[str, Any]: return v.model_dump() diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index f8eb5aa..f0fbbb7 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -16,7 +16,7 @@ ProportionalSetting, ) from phlower.settings._module_settings._share_setting import ShareSetting -from phlower.settings._module_settings._similarity_equivariant_mlp_setting import ( +from phlower.settings._module_settings._similarity_equivariant_mlp_setting import ( # noqa: E501 SimilarityEquivariantMLPSetting, ) diff --git a/src/phlower/settings/_module_settings/_gcn_setting.py b/src/phlower/settings/_module_settings/_gcn_setting.py index 4245771..940c107 100644 --- a/src/phlower/settings/_module_settings/_gcn_setting.py +++ b/src/phlower/settings/_module_settings/_gcn_setting.py @@ -53,7 +53,7 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: @pydantic.model_validator(mode="before") @classmethod - def fill_empty_activations_dropouts(cls, values: dict): + def fill_empty_activations_dropouts(cls, values: dict) -> dict: n_nodes = len(values.get("nodes")) activations = values.get("activations", []) dropouts = values.get("dropouts", []) diff --git a/src/phlower/settings/_module_settings/_mlp_setting.py b/src/phlower/settings/_module_settings/_mlp_setting.py index 20d037a..30cdd0b 100644 --- a/src/phlower/settings/_module_settings/_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_mlp_setting.py @@ -47,7 +47,7 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: @pydantic.model_validator(mode="before") @classmethod - def fill_empty_activations_dropouts(cls, values: dict): + def fill_empty_activations_dropouts(cls, values: dict) -> dict: n_nodes = len(values.get("nodes")) activations = values.get("activations", []) dropouts = values.get("dropouts", []) diff --git a/src/phlower/settings/_phlower_setting.py b/src/phlower/settings/_phlower_setting.py index d926ec2..8453cd0 100644 --- a/src/phlower/settings/_phlower_setting.py +++ b/src/phlower/settings/_phlower_setting.py @@ -123,7 +123,7 @@ class PhlowerPredictorSetting: @pydantic.field_validator("selection_mode") @classmethod - def check_valid_selection_mode(cls, name): + def check_valid_selection_mode(cls, name: str) -> str: names = [v.value for v in ModelSelectionType] if name not in names: raise ValueError(f"{name} selection mode does not exist.") diff --git a/src/phlower/settings/_scaling_setting.py b/src/phlower/settings/_scaling_setting.py index 4980b23..e42970c 100644 --- a/src/phlower/settings/_scaling_setting.py +++ b/src/phlower/settings/_scaling_setting.py @@ -43,14 +43,14 @@ def is_parent_scaler(self) -> bool: @pydantic.field_validator("join_fitting") @classmethod - def must_be_true(cls, v): + def must_be_true(cls, v: bool) -> bool: if not v: raise ValueError("join_fitting must be True except same_as.") return v @pydantic.field_validator("method") @classmethod - def is_exist_method(cls, v): + def is_exist_method(cls, v: str) -> str: logger = utils.get_logger(__name__) names = utils.get_registered_scaler_names() @@ -59,7 +59,7 @@ def is_exist_method(cls, v): logger.warning(f"Scaler name: {v} is not implemented.") return v - def get_scaler_name(self, variable_name: str): + def get_scaler_name(self, variable_name: str) -> str: return f"SCALER_{variable_name}" @@ -76,7 +76,7 @@ class SameAsInputParameters(IScalerParameter, pydantic.BaseModel): def is_parent_scaler(self) -> bool: return False - def get_scaler_name(self, variable_name: str): + def get_scaler_name(self, variable_name: str) -> str: return f"SCALER_{self.same_as}" @@ -92,7 +92,7 @@ class PhlowerScalingSetting(pydantic.BaseModel): model_config = pydantic.ConfigDict(extra="forbid", frozen=True) @pydantic.model_validator(mode="after") - def _check_same_as(self): + def _check_same_as(self) -> Self: for k, v in self.varaible_name_to_scalers.items(): if v.is_parent_scaler: continue @@ -141,7 +141,7 @@ class ScalerResolvedParameter: @pydantic.field_validator("method") @classmethod - def is_exist_method(cls, v): + def is_exist_method(cls, v: str) -> str: logger = utils.get_logger(__name__) names = utils.get_registered_scaler_names() @@ -151,7 +151,7 @@ def is_exist_method(cls, v): return v @pydantic.model_validator(mode="after") - def _validate_isoam_scaler(self): + def _validate_isoam_scaler(self) -> Self: _validate_scaler(self.method, self) return self diff --git a/src/phlower/settings/_trainer_setting.py b/src/phlower/settings/_trainer_setting.py index 4b72a28..370cd7a 100644 --- a/src/phlower/settings/_trainer_setting.py +++ b/src/phlower/settings/_trainer_setting.py @@ -60,7 +60,7 @@ class OptimizerSetting(pydantic.BaseModel): @pydantic.field_validator("optimizer") @classmethod - def check_exist_scheduler(cls, name): + def check_exist_scheduler(cls, name: str) -> str: if not OptimizerSelector.exist(name): raise ValueError( f"{name} is not defined in phlower. " @@ -90,7 +90,7 @@ class SchedulerSetting(pydantic.BaseModel): @pydantic.field_validator("scheduler") @classmethod - def check_exist_scheduler(cls, name): + def check_exist_scheduler(cls, name: str) -> str: if not SchedulerSelector.exist(name): raise ValueError( f"{name} is not defined in phlower. " diff --git a/src/phlower/utils/_logging.py b/src/phlower/utils/_logging.py index ff69ac9..edd96bf 100644 --- a/src/phlower/utils/_logging.py +++ b/src/phlower/utils/_logging.py @@ -5,11 +5,11 @@ class DefaultLoggerFactory: _is_propagate: bool = False @classmethod - def _get_phlower_root_logger(cls): + def _get_phlower_root_logger(cls) -> logging.Logger: return logging.getLogger("phlower") @classmethod - def _get_library_logger(cls): + def _get_library_logger(cls) -> logging.Logger: _phlower_root_logger = logging.getLogger("phlower") _phlower_root_logger.propagate = cls._is_propagate return _phlower_root_logger diff --git a/src/phlower/utils/_optimizer.py b/src/phlower/utils/_optimizer.py index 85b5c73..777aca6 100644 --- a/src/phlower/utils/_optimizer.py +++ b/src/phlower/utils/_optimizer.py @@ -22,7 +22,7 @@ def register(name: str, cls: type[torch.optim.Optimizer]) -> None: OptimizerSelector._REGISTERED.update({name: cls}) @staticmethod - def exist(name: str): + def exist(name: str) -> bool: return name in OptimizerSelector._REGISTERED @staticmethod diff --git a/src/phlower/utils/_schedulers.py b/src/phlower/utils/_schedulers.py index 097a8ae..99ca30a 100644 --- a/src/phlower/utils/_schedulers.py +++ b/src/phlower/utils/_schedulers.py @@ -27,7 +27,7 @@ def register( SchedulerSelector._REGISTERED.update({name: cls}) @staticmethod - def exist(name: str): + def exist(name: str) -> bool: return name in SchedulerSelector._REGISTERED @staticmethod diff --git a/src/phlower/utils/preprocess.py b/src/phlower/utils/preprocess.py index 12cf2c4..b6cd52a 100644 --- a/src/phlower/utils/preprocess.py +++ b/src/phlower/utils/preprocess.py @@ -10,7 +10,7 @@ def get_registered_scaler_names() -> list[str]: return [v.value for v in PhlowerScalerName] -def convert_to_dumped(v: Any): +def convert_to_dumped(v: Any) -> Any: # noqa: ANN401 if isinstance(v, DenseArrayType): return v.tolist() diff --git a/tests/e2e_tests/test_preprocess.py b/tests/e2e_tests/test_preprocess.py index 9231eb3..eda0d5d 100644 --- a/tests/e2e_tests/test_preprocess.py +++ b/tests/e2e_tests/test_preprocess.py @@ -4,11 +4,10 @@ import numpy as np import pytest import scipy.sparse as sp -from pipe import select, where - from phlower.io import PhlowerDirectory, PhlowerFileBuilder from phlower.services.preprocessing import PhlowerScalingService from phlower.settings import PhlowerScalingSetting, PhlowerSetting +from pipe import select, where _OUTPUT_DIR = pathlib.Path(__file__).parent / "_tmp_preprocess" @@ -55,7 +54,7 @@ def prepare_sample_interim_files(): @pytest.fixture(scope="module") -def perform_scaling(prepare_sample_interim_files): +def perform_scaling(prepare_sample_interim_files: None) -> PhlowerSetting: path = _OUTPUT_DIR phlower_path = PhlowerDirectory(path) @@ -83,7 +82,9 @@ def perform_scaling(prepare_sample_interim_files): [(_OUTPUT_DIR / "interim", _OUTPUT_DIR / "preprocessed")], ) def test__saved_array_is_same_as_saved_scalers_transformed( - interim_base_directory, scaling_base_directory, perform_scaling + interim_base_directory: pathlib.Path, + scaling_base_directory: pathlib.Path, + perform_scaling: None, ): phlower_path = PhlowerDirectory(interim_base_directory) interim_directories = list( diff --git a/tests/e2e_tests/test_train.py b/tests/e2e_tests/test_train.py index a94ac5b..7771f9e 100644 --- a/tests/e2e_tests/test_train.py +++ b/tests/e2e_tests/test_train.py @@ -6,7 +6,6 @@ import pytest import scipy.sparse as sp import torch - from phlower import PhlowerTensor from phlower.io import PhlowerDirectory from phlower.services.predictor import PhlowerPredictor @@ -55,7 +54,7 @@ def prepare_sample_preprocessed_files(): @pytest.fixture(scope="module") -def simple_training(prepare_sample_preprocessed_files): +def simple_training(prepare_sample_preprocessed_files: None) -> PhlowerTensor: phlower_path = PhlowerDirectory(_OUTPUT_DIR) preprocessed_directories = list( @@ -80,7 +79,9 @@ def simple_training(prepare_sample_preprocessed_files): @pytest.mark.e2e_test -def test__training_with_multiple_batch_size(prepare_sample_preprocessed_files): +def test__training_with_multiple_batch_size( + prepare_sample_preprocessed_files: None, +): phlower_path = PhlowerDirectory(_OUTPUT_DIR) preprocessed_directories = list( @@ -110,7 +111,7 @@ def test__training_with_multiple_batch_size(prepare_sample_preprocessed_files): @pytest.mark.e2e_test -def test__simple_training(simple_training): +def test__simple_training(simple_training: PhlowerTensor): loss: PhlowerTensor = simple_training assert loss.has_dimension @@ -119,7 +120,7 @@ def test__simple_training(simple_training): @pytest.mark.e2e_test -def test__predict(simple_training): +def test__predict(simple_training: PhlowerTensor): setting = PhlowerSetting.read_yaml("tests/e2e_tests/data/predict.yml") model_directory = _OUTPUT_DIR / "model" diff --git a/tests/test_base/test_array/test_sparse_array_wrapper.py b/tests/test_base/test_array/test_sparse_array_wrapper.py index 951a288..e20d39e 100644 --- a/tests/test_base/test_array/test_sparse_array_wrapper.py +++ b/tests/test_base/test_array/test_sparse_array_wrapper.py @@ -1,7 +1,6 @@ import numpy as np import pytest import scipy.sparse as sp - from phlower._base.array.sparse import ( SparseArrayWrapper, batch, @@ -17,7 +16,7 @@ ([(3, 5)], (3, 5)), ], ) -def test__batch(shapes, expected_shape): +def test__batch(shapes: tuple[int], expected_shape: tuple[int]): rng = np.random.default_rng() sparse_arrays = [ @@ -35,7 +34,7 @@ def test__batch(shapes, expected_shape): "shapes", [([(5, 6), (4, 9), (10, 11)]), ([(1, 1), (2, 1), (1, 1)]), ([(3, 5)])], ) -def test__unbatch(shapes): +def test__unbatch(shapes: list[tuple[int]]): rng = np.random.default_rng() sparse_arrays = [ diff --git a/tests/test_base/test_dimensions.py b/tests/test_base/test_dimensions.py index 262ebd3..3f3e547 100644 --- a/tests/test_base/test_dimensions.py +++ b/tests/test_base/test_dimensions.py @@ -1,7 +1,6 @@ import pytest from hypothesis import assume, given from hypothesis import strategies as st - from phlower._base import PhysicalDimensions from phlower.utils.enums import PhysicalDimensionSymbolType from phlower.utils.exceptions import InvalidDimensionError @@ -13,7 +12,7 @@ st.floats(allow_nan=False), ) ) -def test__equal_when_same_dimension(dict_data): +def test__equal_when_same_dimension(dict_data: dict[str.float]): dimension = PhysicalDimensions(dict_data) other = PhysicalDimensions(dict_data) @@ -24,7 +23,7 @@ def test__equal_when_same_dimension(dict_data): "dict_data", [({"kg": 2, "mm": 3}), ({"m": 2, "hour": 3.2}), ({"mass": None})], ) -def test__failed_when_not_exist_key(dict_data): +def test__failed_when_not_exist_key(dict_data: dict[str, float]): with pytest.raises(InvalidDimensionError): _ = PhysicalDimensions(dict_data) @@ -45,7 +44,7 @@ def test__failed_when_not_exist_key(dict_data): ), ) ) -def test__not_equal_dimension(tuple_dict_data): +def test__not_equal_dimension(tuple_dict_data: tuple[dict, dict]): dict_data1, dict_data2 = tuple_dict_data assume(dict_data1 != dict_data2) @@ -68,7 +67,7 @@ def test__default_dimension(): st.floats(allow_nan=False), ) ) -def test__to_list(dict_data): +def test__to_list(dict_data: dict[str, float]): dimension = PhysicalDimensions(dict_data) list_data = dimension.to_list() diff --git a/tests/test_base/test_functionals/conftest.py b/tests/test_base/test_functionals/conftest.py index 8a3f238..26138e4 100644 --- a/tests/test_base/test_functionals/conftest.py +++ b/tests/test_base/test_functionals/conftest.py @@ -2,10 +2,9 @@ import numpy as np import pytest -from scipy import sparse as sp - from phlower._base.array import phlower_array from phlower._base.tensors._interface import IPhlowerTensor +from scipy import sparse as sp @pytest.fixture diff --git a/tests/test_base/test_functionals/test_batch.py b/tests/test_base/test_functionals/test_batch.py index 9f657dc..d8be3bf 100644 --- a/tests/test_base/test_functionals/test_batch.py +++ b/tests/test_base/test_functionals/test_batch.py @@ -1,7 +1,9 @@ -import pytest +from collections.abc import Callable +import pytest from phlower._base._functionals import to_batch from phlower._base.tensors import phlower_dimension_tensor +from phlower._base.tensors._interface import IPhlowerTensor @pytest.mark.parametrize( @@ -20,7 +22,12 @@ ], ) def test__to_batch_for_sparse_tensors( - shapes, dimensions, desired_shape, create_sparse_tensors + shapes: list[tuple[int]], + dimensions: dict | None, + desired_shape: tuple[int], + create_sparse_tensors: Callable[ + [list[tuple[int]], list[dict[str, float]] | None], list[IPhlowerTensor] + ], ): tensors = create_sparse_tensors(shapes, dimensions) batched_tensor, batch_info = to_batch(tensors) @@ -54,7 +61,13 @@ def test__to_batch_for_sparse_tensors( ], ) def test__to_batch_for_dense_tensors( - shapes, concat_dim, dimensions, desired_shape, create_dense_tensors + shapes: list[tuple[int]], + concat_dim: int, + dimensions: dict, + desired_shape: tuple[int], + create_dense_tensors: Callable[ + [list[tuple[int]], list[dict[str, float]] | None], list[IPhlowerTensor] + ], ): tensors = create_dense_tensors(shapes, dimensions) batched_tensor, batch_info = to_batch(tensors, concat_dim) diff --git a/tests/test_base/test_functionals/test_check.py b/tests/test_base/test_functionals/test_check.py index e0eab88..d2ef15e 100644 --- a/tests/test_base/test_functionals/test_check.py +++ b/tests/test_base/test_functionals/test_check.py @@ -1,12 +1,14 @@ import numpy as np import pytest import torch - from phlower._base import phlower_dimension_tensor, phlower_tensor from phlower._base._functionals import is_same_dimensions, is_same_layout +from phlower._base.tensors._interface import IPhlowerTensor -def _create_random_tensors(shapes: list[tuple], is_sparse: list[bool]): +def _create_random_tensors( + shapes: list[tuple], is_sparse: list[bool] +) -> list[IPhlowerTensor]: assert len(shapes) == len(is_sparse) tensors = [torch.tensor(np.random.rand(*shape)) for shape in shapes] tensors = [ @@ -17,7 +19,7 @@ def _create_random_tensors(shapes: list[tuple], is_sparse: list[bool]): def _create_random_tensors_with_dims( shapes: list[tuple], dimensions: list[dict[str, float]] -): +) -> list[IPhlowerTensor]: assert len(shapes) == len(dimensions) tensors = [torch.tensor(np.random.rand(*shape)) for shape in shapes] _dimensions = [phlower_dimension_tensor(v) for v in dimensions] @@ -33,7 +35,9 @@ def _create_random_tensors_with_dims( ([(3, 6), (11, 5)], [False, False], True), ], ) -def test__is_same_layout(shapes, is_sparse, desired): +def test__is_same_layout( + shapes: list[tuple[int]], is_sparse: list[bool], desired: bool +): tensors = _create_random_tensors(shapes, is_sparse) assert is_same_layout(tensors) == desired @@ -70,6 +74,8 @@ def test__is_same_layout(shapes, is_sparse, desired): ), ], ) -def test__is_same_dimensions(shapes, dimensions, desired): +def test__is_same_dimensions( + shapes: list[tuple[int]], dimensions: list[dict], desired: bool +): tensors = _create_random_tensors_with_dims(shapes, dimensions) assert is_same_dimensions(tensors) == desired diff --git a/tests/test_base/test_functionals/test_unbatch.py b/tests/test_base/test_functionals/test_unbatch.py index b44f7e3..c46301c 100644 --- a/tests/test_base/test_functionals/test_unbatch.py +++ b/tests/test_base/test_functionals/test_unbatch.py @@ -1,8 +1,10 @@ +from collections.abc import Callable + import numpy as np import pytest import torch - from phlower._base._functionals import to_batch, unbatch +from phlower._base.tensors._interface import IPhlowerTensor @pytest.mark.parametrize( @@ -23,7 +25,12 @@ ], ) def test__unbatch_for_sparse( - shapes, dimensions, expected_shape, create_sparse_tensors + shapes: list[tuple[int]], + dimensions: list[dict], + expected_shape: tuple, + create_sparse_tensors: Callable[ + [list[tuple[int]], list[dict[str, float]] | None], list[IPhlowerTensor] + ], ): sparse_tensors = create_sparse_tensors(shapes, dimensions) @@ -62,7 +69,13 @@ def test__unbatch_for_sparse( ], ) def test__unbatch_for_dense( - shapes, concat_dim, dimensions, expected_shape, create_dense_tensors + shapes: list[tuple[int]], + concat_dim: int, + dimensions: dict | None, + expected_shape: tuple[int], + create_dense_tensors: Callable[ + [list[tuple[int]], list[dict[str, float]] | None], list[IPhlowerTensor] + ], ): dense_tensors = create_dense_tensors(shapes, dimensions) @@ -92,11 +105,15 @@ def test__unbatch_for_dense( ], ) def test__batched_tensor_mm( - sparse_shapes, - dense_shapes, - expected_shape, - create_sparse_tensors, - create_dense_tensors, + sparse_shapes: list[tuple[int]], + dense_shapes: list[tuple[int]], + expected_shape: tuple[int], + create_sparse_tensors: Callable[ + [list[tuple[int]], list[dict[str, float]] | None], list[IPhlowerTensor] + ], + create_dense_tensors: Callable[ + [list[tuple[int]], list[dict[str, float]] | None], list[IPhlowerTensor] + ], ): sparse_tensors = create_sparse_tensors(sparse_shapes) dense_tensors = create_dense_tensors(dense_shapes) diff --git a/tests/test_base/test_tensors/test__dimensions.py b/tests/test_base/test_tensors/test__dimensions.py index cd7f0a2..cc09106 100644 --- a/tests/test_base/test_tensors/test__dimensions.py +++ b/tests/test_base/test_tensors/test__dimensions.py @@ -1,6 +1,5 @@ import pytest import torch - from phlower import PhlowerDimensionTensor from phlower.utils.exceptions import DimensionIncompatibleError @@ -8,7 +7,7 @@ @pytest.mark.parametrize( "inputs", [[0, 2, 0, 0, 0, 0, 0], [0, 2, 0, 2, 0, 0, 0]] ) -def test__initialize(inputs): +def test__initialize(inputs: list[int]): _ = PhlowerDimensionTensor.from_list(inputs) @@ -25,7 +24,7 @@ def test__initialize(inputs): ), ], ) -def test__add(unit1, unit2): +def test__add(unit1: list[int], unit2: list[int]): unit1 = PhlowerDimensionTensor.from_list(unit1) unit2 = PhlowerDimensionTensor.from_list(unit2) @@ -40,7 +39,7 @@ def test__add(unit1, unit2): ], ) @pytest.mark.parametrize("dim", [0, 2]) -def test__cat(unit, dim): +def test__cat(unit: list[int], dim: int): unit1 = PhlowerDimensionTensor.from_list(unit) unit2 = PhlowerDimensionTensor.from_list(unit) @@ -60,7 +59,7 @@ def test__cat(unit, dim): ), ], ) -def test__cat_raise_dimension_incompatible(unit1, unit2): +def test__cat_raise_dimension_incompatible(unit1: list[int], unit2: list[int]): unit1 = PhlowerDimensionTensor.from_list(unit1) unit2 = PhlowerDimensionTensor.from_list(unit2) diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index e931701..360cf70 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -1,7 +1,6 @@ import numpy as np import pytest import torch - from phlower import PhlowerTensor, phlower_dimension_tensor, phlower_tensor from phlower.utils.exceptions import ( DimensionIncompatibleError, @@ -51,7 +50,9 @@ def test__sub_with_unit(): @pytest.mark.parametrize( "unit1, unit2", [({"L": 2, "T": -2}, None), (None, {"M": 2, "T": -3})] ) -def test__add_with_and_without_dimensions(unit1, unit2): +def test__add_with_and_without_dimensions( + unit1: dict | None, unit2: dict | None +): tensor1 = phlower_tensor(torch.rand(3, 4), dimension=unit1) tensor2 = phlower_tensor(torch.rand(3, 4), dimension=unit2) @@ -163,7 +164,9 @@ def test__tanh(): (True, True, [4, 10, 10, 10, 3, 3, 16], 2), ], ) -def test__rank(is_time_series, is_voxel, size, desired_rank): +def test__rank( + is_time_series: bool, is_voxel: bool, size: list[int], desired_rank: int +): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel @@ -188,7 +191,12 @@ def test__rank(is_time_series, is_voxel, size, desired_rank): (True, True, [4, 10, 10, 10, 3, 3, 16], 1000), ], ) -def test__n_vertices(is_time_series, is_voxel, size, desired_n_vertices): +def test__n_vertices( + is_time_series: bool, + is_voxel: bool, + size: list[int], + desired_n_vertices: int, +): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel @@ -220,7 +228,12 @@ def test__raises_phlower_sparse_rank_undefined_error(): (True, True, [4, 10, 10, 10, 3, 3, 16], (1000, 4 * 3 * 3 * 16)), ], ) -def test__to_vertexwise(is_time_series, is_voxel, size, desired_shape): +def test__to_vertexwise( + is_time_series: bool, + is_voxel: bool, + size: list[int], + desired_shape: tuple[int], +): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel @@ -245,7 +258,9 @@ def test__to_vertexwise(is_time_series, is_voxel, size, desired_shape): (True, True, [4, 10, 10, 10, 3, 3, 16]), ], ) -def test__to_vertexwise_inverse(is_time_series, is_voxel, size): +def test__to_vertexwise_inverse( + is_time_series: bool, is_voxel: bool, size: list[int] +): torch_tensor = torch.rand(*size) phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel @@ -270,7 +285,12 @@ def test__to_vertexwise_inverse(is_time_series, is_voxel, size): ((10, 3 * 16), "n (p a) -> n p a", {"p": 3}, (10, 3, 16)), ], ) -def test__rearrange(input_shape, pattern, dict_shape, desired_shape): +def test__rearrange( + input_shape: tuple[int], + pattern: str, + dict_shape: dict, + desired_shape: tuple[int], +): phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) actual = phlower_tensor.rearrange(pattern, **dict_shape) assert actual.shape == desired_shape diff --git a/tests/test_data/conftest.py b/tests/test_data/conftest.py index 3a37a52..861abc5 100644 --- a/tests/test_data/conftest.py +++ b/tests/test_data/conftest.py @@ -5,7 +5,6 @@ import numpy as np import pytest import scipy.sparse as sp - from phlower.io import PhlowerNumpyFile from phlower.utils.typing import ArrayDataType @@ -13,12 +12,12 @@ @pytest.fixture -def output_base_directory(): +def output_base_directory() -> pathlib.Path: return _output_base_directory @pytest.fixture(scope="module") -def create_tmp_dataset(): +def create_tmp_dataset() -> dict[str, dict[str, ArrayDataType]]: if _output_base_directory.exists(): shutil.rmtree(_output_base_directory) _output_base_directory.mkdir(parents=True) diff --git a/tests/test_data/test_data_loader.py b/tests/test_data/test_data_loader.py index 08d1d18..4d7e8c0 100644 --- a/tests/test_data/test_data_loader.py +++ b/tests/test_data/test_data_loader.py @@ -1,7 +1,9 @@ +import pathlib +from collections.abc import Callable + import pytest from hypothesis import given from hypothesis import strategies as st - from phlower._base import PhysicalDimensions from phlower.data import ( DataLoaderBuilder, @@ -13,7 +15,7 @@ @st.composite -def trainer_setting(draw): +def trainer_setting(draw: Callable) -> PhlowerTrainerSetting: setting = PhlowerTrainerSetting( loss_setting={"name2loss": {"u": "mse"}}, non_blocking=draw(st.booleans()), @@ -26,7 +28,7 @@ def trainer_setting(draw): @st.composite -def predictor_setting(draw): +def predictor_setting(draw: Callable) -> PhlowerPredictorSetting: setting = PhlowerPredictorSetting( selection_mode=draw(st.sampled_from(ModelSelectionType)), non_blocking=draw(st.booleans()), @@ -61,7 +63,9 @@ def test__create_from_predictor_setting(setting: PhlowerPredictorSetting): @pytest.mark.parametrize("batch_size", [1, 2, 3]) def test__consider_batch_size( - batch_size, create_tmp_dataset, output_base_directory + batch_size: int, + create_tmp_dataset: None, + output_base_directory: pathlib.Path, ): directories = [ output_base_directory / v for v in ["data0", "data1", "data2"] @@ -110,11 +114,11 @@ def test__consider_batch_size( ], ) def test__consider_dimensions( - dimensions, - disable_dimensions, - desired, - create_tmp_dataset, - output_base_directory, + dimensions: dict, + disable_dimensions: bool, + desired: dict, + create_tmp_dataset: None, + output_base_directory: pathlib.Path, ): directories = [ output_base_directory / v for v in ["data0", "data1", "data2"] @@ -172,7 +176,10 @@ def test__consider_dimensions( ], ) def test__not_consider_dimensions( - dimensions, disable_dimensions, create_tmp_dataset, output_base_directory + dimensions: dict, + disable_dimensions: bool, + create_tmp_dataset: None, + output_base_directory: pathlib.Path, ): directories = [ output_base_directory / v for v in ["data0", "data1", "data2"] diff --git a/tests/test_data/test_datasets.py b/tests/test_data/test_datasets.py index 9798fdb..d051441 100644 --- a/tests/test_data/test_datasets.py +++ b/tests/test_data/test_datasets.py @@ -1,6 +1,7 @@ +import pathlib + import numpy as np import pytest - from phlower.data import LazyPhlowerDataset from phlower.utils.typing import ArrayDataType @@ -10,7 +11,10 @@ [(["data0", "data1"], 2), (["data0", "data1", "data2"], 3)], ) def test__lazy_dataset_length( - directories, desired, create_tmp_dataset, output_base_directory + directories: list[str], + desired: int, + create_tmp_dataset: None, + output_base_directory: pathlib.Path, ): directories = [output_base_directory / v for v in directories] dataset = LazyPhlowerDataset( @@ -27,12 +31,12 @@ def test__lazy_dataset_length( [(["x0", "x1", "x2"], ["y0"], ["s0", "s1"], ["data0", "data1", "data2"])], ) def test__lazy_dataset_getitem( - x_variable_names, - y_variable_names, - support_names, - directory_names, - create_tmp_dataset, - output_base_directory, + x_variable_names: list[str], + y_variable_names: list[str], + support_names: list[str], + directory_names: list[str], + create_tmp_dataset: None, + output_base_directory: pathlib.Path, ): directories = [output_base_directory / v for v in directory_names] dataset = LazyPhlowerDataset( @@ -75,12 +79,12 @@ def test__lazy_dataset_getitem( ], ) def test__lazy_dataset_getitem_when_no_ydata( - x_variable_names, - y_variable_names, - support_names, - directory_names, - create_tmp_dataset, - output_base_directory, + x_variable_names: list[str], + y_variable_names: list[str], + support_names: list[str], + directory_names: list[str], + create_tmp_dataset: None, + output_base_directory: pathlib.Path, ): directories = [output_base_directory / v for v in directory_names] dataset = LazyPhlowerDataset( diff --git a/tests/test_io/test_directory.py b/tests/test_io/test_directory.py index bca41b1..c6e2463 100644 --- a/tests/test_io/test_directory.py +++ b/tests/test_io/test_directory.py @@ -2,7 +2,6 @@ import shutil import pytest - from phlower.io import PhlowerDirectory TEST_DATA_DIR = pathlib.Path(__file__).parent / "tmp" @@ -21,7 +20,7 @@ def create_test_cases(): @pytest.mark.parametrize( "path", ["tests/data/deform", "tests/data/csv_prepost/raw"] ) -def test__initialize(path): +def test__initialize(path: str): phlower_dir = PhlowerDirectory(pathlib.Path(path)) assert phlower_dir.path == pathlib.Path(path) @@ -30,7 +29,9 @@ def test__initialize(path): @pytest.mark.parametrize( "variable_name, ext", [("variable1", ".npy"), ("variable2", ".npz.enc")] ) -def test__find_variable_file(variable_name, ext, create_test_cases): +def test__find_variable_file( + variable_name: str, ext: str, create_test_cases: None +): phlower_dir = PhlowerDirectory(TEST_DATA_DIR) phlower_file = phlower_dir.find_variable_file(variable_name) @@ -41,7 +42,9 @@ def test__find_variable_file(variable_name, ext, create_test_cases): @pytest.mark.parametrize( "variable_name, ext", [("variable3", ".npz"), ("variable4", ".npz.enc")] ) -def test__failed_find_variable_file(variable_name, ext, create_test_cases): +def test__failed_find_variable_file( + variable_name: str, ext: str, create_test_cases: None +): phlower_dir = PhlowerDirectory(TEST_DATA_DIR) assert phlower_dir.exist_variable_file(variable_name) is False @@ -52,7 +55,9 @@ def test__failed_find_variable_file(variable_name, ext, create_test_cases): @pytest.mark.parametrize( "variable_name, ext", [("variable3", ".npz"), ("variable4", ".npz.enc")] ) -def test__find_variable_file_as_None(variable_name, ext, create_test_cases): +def test__find_variable_file_as_None( + variable_name: str, ext: str, create_test_cases: None +): phlower_dir = PhlowerDirectory(TEST_DATA_DIR) phlower_file = phlower_dir.find_variable_file( diff --git a/tests/test_io/test_files/conftest.py b/tests/test_io/test_files/conftest.py index af11941..b2dc35b 100644 --- a/tests/test_io/test_files/conftest.py +++ b/tests/test_io/test_files/conftest.py @@ -5,7 +5,7 @@ @pytest.fixture(scope="module") -def setup_test_dir(): +def setup_test_dir() -> pathlib.Path: directory = pathlib.Path(__file__).parent test_dir = directory / "tmp" if test_dir.exists(): diff --git a/tests/test_io/test_files/test_checkpoint_file.py b/tests/test_io/test_files/test_checkpoint_file.py index 8cce885..ffe8959 100644 --- a/tests/test_io/test_files/test_checkpoint_file.py +++ b/tests/test_io/test_files/test_checkpoint_file.py @@ -4,7 +4,6 @@ import numpy as np import pytest import torch - from phlower.io import PhlowerCheckpointFile TEST_ENCRYPT_KEY = secrets.token_bytes(32) @@ -14,7 +13,7 @@ "path, ext", [("./sample/sample.pth", ".pth"), ("./sample/sample.pth.enc", ".pth.enc")], ) -def test__check_extension_type(path, ext): +def test__check_extension_type(path: str, ext: str): path = pathlib.Path(path) siml_path = PhlowerCheckpointFile(path) assert siml_path._ext_type.value == ext @@ -24,7 +23,7 @@ def test__check_extension_type(path, ext): @pytest.mark.parametrize( "path", [("./sample/sample.npy"), ("./sample/sample.npy.enc")] ) -def test__check_error_extension_type(path): +def test__check_error_extension_type(path: str): path = pathlib.Path(path) with pytest.raises(NotImplementedError): _ = PhlowerCheckpointFile(path) @@ -34,13 +33,13 @@ def test__check_error_extension_type(path): "path, enc", [("./sample/sample.pth", False), ("./sample/sample.pth.enc", True)], ) -def test__is_encrypted(path, enc): +def test__is_encrypted(path: str, enc: bool): path = pathlib.Path(path) siml_path = PhlowerCheckpointFile(path) assert siml_path.is_encrypted == enc -def test__save_and_load(setup_test_dir): +def test__save_and_load(setup_test_dir: pathlib.Path): sample_tensor = torch.tensor(np.random.rand(3, 4)) saved_path = PhlowerCheckpointFile.save( @@ -57,7 +56,7 @@ def test__save_and_load(setup_test_dir): np.testing.assert_array_almost_equal(loaded_data, sample_tensor) -def test__save_encrypted_and_load(setup_test_dir): +def test__save_encrypted_and_load(setup_test_dir: pathlib.Path): sample_tensor = torch.tensor(np.random.rand(3, 4)) output_directory = setup_test_dir @@ -75,7 +74,7 @@ def test__save_encrypted_and_load(setup_test_dir): np.testing.assert_array_almost_equal(loaded_data, sample_tensor) -def test__cannnot_load_without_key(setup_test_dir): +def test__cannnot_load_without_key(setup_test_dir: pathlib.Path): sample_tensor = torch.tensor(np.random.rand(3, 4)) output_directory = setup_test_dir @@ -93,7 +92,7 @@ def test__cannnot_load_without_key(setup_test_dir): _ = saved_path.load(device="cpu") -def test__save_not_allowed_overwrite(setup_test_dir): +def test__save_not_allowed_overwrite(setup_test_dir: pathlib.Path): sample_array = np.random.rand(3, 4) sample_tensor = torch.tensor(sample_array) @@ -117,7 +116,7 @@ def test__save_not_allowed_overwrite(setup_test_dir): ("./aaa/snapshot_epoch_30.pth.enc", 30), ], ) -def test__get_epoch(path, num_epoch): +def test__get_epoch(path: str, num_epoch: int): phlower_path = PhlowerCheckpointFile(path) assert phlower_path.epoch == num_epoch @@ -126,7 +125,7 @@ def test__get_epoch(path, num_epoch): "path", [("./aaa/deployed_1.pth"), ("./aaa/model.pth"), ("./aaa/epoch_2.pth")], ) -def test__get_epoch_not_handled(path): +def test__get_epoch_not_handled(path: str): phlower_path = PhlowerCheckpointFile(path) with pytest.raises(ValueError): diff --git a/tests/test_io/test_files/test_numpy_file.py b/tests/test_io/test_files/test_numpy_file.py index 80b87b8..077fe04 100644 --- a/tests/test_io/test_files/test_numpy_file.py +++ b/tests/test_io/test_files/test_numpy_file.py @@ -4,8 +4,8 @@ import numpy as np import pytest import scipy.sparse as sp - from phlower.io import PhlowerNumpyFile +from phlower.utils.typing import ArrayDataType TEST_ENCRYPT_KEY = secrets.token_bytes(32) @@ -19,7 +19,7 @@ ("./sample/sample.npz.enc", ".npz.enc"), ], ) -def test__check_extension_type(path, ext): +def test__check_extension_type(path: str, ext: str): path = pathlib.Path(path) phlower_path = PhlowerNumpyFile(path) assert phlower_path._ext_type.value == ext @@ -29,7 +29,7 @@ def test__check_extension_type(path, ext): @pytest.mark.parametrize( "path", [("./sample/sample.pkl"), ("./sample/sample.pkl.enc")] ) -def test__check_error_extension_type(path): +def test__check_error_extension_type(path: str): path = pathlib.Path(path) with pytest.raises(NotImplementedError): _ = PhlowerNumpyFile(path) @@ -44,7 +44,7 @@ def test__check_error_extension_type(path): ("./sample/sample.npz.enc", True), ], ) -def test__is_encrypted(path, enc): +def test__is_encrypted(path: str, enc: bool): path = pathlib.Path(path) phlower_path = PhlowerNumpyFile(path) assert phlower_path.is_encrypted == enc @@ -58,7 +58,11 @@ def test__is_encrypted(path, enc): (TEST_ENCRYPT_KEY, TEST_ENCRYPT_KEY), ], ) -def test__save_npy_and_load(encrypt_key, decrypt_key, setup_test_dir): +def test__save_npy_and_load( + encrypt_key: bytes | None, + decrypt_key: bytes | None, + setup_test_dir: pathlib.Path, +): sample_array = np.random.rand(3, 4) saved_path = PhlowerNumpyFile.save( @@ -83,7 +87,11 @@ def test__save_npy_and_load(encrypt_key, decrypt_key, setup_test_dir): (TEST_ENCRYPT_KEY, TEST_ENCRYPT_KEY), ], ) -def test__save_npz_and_load(encrypt_key, decrypt_key, setup_test_dir): +def test__save_npz_and_load( + encrypt_key: bytes | None, + decrypt_key: bytes | None, + setup_test_dir: pathlib.Path, +): rng = np.random.default_rng() sample_array = sp.random(5, 5, density=0.1, random_state=rng) @@ -106,7 +114,9 @@ def test__save_npz_and_load(encrypt_key, decrypt_key, setup_test_dir): @pytest.mark.parametrize( "data", [np.random.rand(3, 5), sp.random(5, 5, density=0.1)] ) -def test__cannnot_load_encrypt_data_without_key(data, setup_test_dir): +def test__cannnot_load_encrypt_data_without_key( + data: list[ArrayDataType], setup_test_dir: pathlib.Path +): saved_path = PhlowerNumpyFile.save( output_directory=setup_test_dir, file_basename="sample", @@ -119,7 +129,7 @@ def test__cannnot_load_encrypt_data_without_key(data, setup_test_dir): _ = saved_path.load() -def test__save_not_allowed_overwrite(setup_test_dir): +def test__save_not_allowed_overwrite(setup_test_dir: pathlib.Path): sample_array = np.random.rand(3, 4) path: pathlib.Path = setup_test_dir / "sample.npy" diff --git a/tests/test_io/test_files/test_yaml_file.py b/tests/test_io/test_files/test_yaml_file.py index 39b3296..a0d8b4d 100644 --- a/tests/test_io/test_files/test_yaml_file.py +++ b/tests/test_io/test_files/test_yaml_file.py @@ -2,7 +2,6 @@ import secrets import pytest - from phlower.io import PhlowerYamlFile TEST_ENCRYPT_KEY = secrets.token_bytes(32) @@ -12,7 +11,7 @@ "path, ext", [("./sample/sample.yml", ".yml"), ("./sample/sample.yml.enc", ".yml.enc")], ) -def test__check_extension_type(path, ext): +def test__check_extension_type(path: str, ext: str): path = pathlib.Path(path) phlower_path = PhlowerYamlFile(path) assert phlower_path._ext_type.value == ext @@ -22,7 +21,7 @@ def test__check_extension_type(path, ext): @pytest.mark.parametrize( "path", [("./sample/sample.npy"), ("./sample/sample.npy.enc")] ) -def test__check_error_extension_type(path): +def test__check_error_extension_type(path: str): path = pathlib.Path(path) with pytest.raises(NotImplementedError): _ = PhlowerYamlFile(path) @@ -32,7 +31,7 @@ def test__check_error_extension_type(path): "path, enc", [("./sample/sample.yml", False), ("./sample/sample.yml.enc", True)], ) -def test__is_encrypted(path, enc): +def test__is_encrypted(path: str, enc: bool): path = pathlib.Path(path) phlower_path = PhlowerYamlFile(path) assert phlower_path.is_encrypted == enc @@ -46,7 +45,11 @@ def test__is_encrypted(path, enc): (TEST_ENCRYPT_KEY, TEST_ENCRYPT_KEY), ], ) -def test__save_and_load(encrypt_key, decrypt_key, setup_test_dir): +def test__save_and_load( + encrypt_key: bytes | None, + decrypt_key: bytes | None, + setup_test_dir: pathlib.Path, +): sample_data = {"a": 1, "b": 2} saved_path = PhlowerYamlFile.save( @@ -63,7 +66,7 @@ def test__save_and_load(encrypt_key, decrypt_key, setup_test_dir): assert loaded_data == sample_data -def test__cannnot_load_encrypt_data_without_key(setup_test_dir): +def test__cannnot_load_encrypt_data_without_key(setup_test_dir: pathlib.Path): sample_data = {"a": 1, "b": 2} saved_path = PhlowerYamlFile.save( @@ -78,7 +81,7 @@ def test__cannnot_load_encrypt_data_without_key(setup_test_dir): _ = saved_path.load() -def test__save_not_allowed_overwrite(setup_test_dir): +def test__save_not_allowed_overwrite(setup_test_dir: pathlib.Path): sample_data = {"a": 1, "b": 2} path = setup_test_dir / "sample.yml" diff --git a/tests/test_io/test_model_selector.py b/tests/test_io/test_model_selector.py index 71fbfc7..0080b07 100644 --- a/tests/test_io/test_model_selector.py +++ b/tests/test_io/test_model_selector.py @@ -6,7 +6,6 @@ import pytest from hypothesis import given from hypothesis import strategies as st - from phlower.io import select_snapshot_file from phlower.io._model_selector import ModelSelectorBuilder from phlower.utils.enums import ModelSelectionType @@ -15,12 +14,12 @@ @given(st.sampled_from(ModelSelectionType)) -def test__selctor_builder(select_type): +def test__selctor_builder(select_type: ModelSelectionType): _ = ModelSelectorBuilder.create(select_type.value) @pytest.mark.parametrize("selection_name", ["best_of_best", "none"]) -def test__not_implemented_selection_name(selection_name): +def test__not_implemented_selection_name(selection_name: str): with pytest.raises(NotImplementedError): _ = ModelSelectorBuilder.create(selection_name) @@ -51,7 +50,7 @@ def prepare_snapshots(): df.to_csv(TEST_DATA_DIR / "log.csv") -def test__best_select_model(prepare_snapshots): +def test__best_select_model(prepare_snapshots: None): actual_path = select_snapshot_file( TEST_DATA_DIR, selection_mode=ModelSelectionType.BEST.value ) @@ -63,7 +62,7 @@ def test__best_select_model(prepare_snapshots): assert actual_path.epoch == epoch -def test__latest_select_model(prepare_snapshots): +def test__latest_select_model(prepare_snapshots: None): actual_path = select_snapshot_file( TEST_DATA_DIR, selection_mode=ModelSelectionType.LATEST.value ) @@ -73,7 +72,7 @@ def test__latest_select_model(prepare_snapshots): assert actual_path.epoch == max_epoch -def test__train_best_select_model(prepare_snapshots): +def test__train_best_select_model(prepare_snapshots: None): actual_path = select_snapshot_file( TEST_DATA_DIR, selection_mode=ModelSelectionType.TRAIN_BEST.value ) @@ -86,7 +85,7 @@ def test__train_best_select_model(prepare_snapshots): @pytest.mark.parametrize("epoch", [1, 5, 6, 8]) -def test__spcified_model_selector(epoch, prepare_snapshots): +def test__spcified_model_selector(epoch: int, prepare_snapshots: None): actual_path = select_snapshot_file( TEST_DATA_DIR, selection_mode=ModelSelectionType.SPECIFIED.value, @@ -97,7 +96,9 @@ def test__spcified_model_selector(epoch, prepare_snapshots): @pytest.mark.parametrize("epoch", [100, 200]) -def test__spcified_model_selector_not_existed(epoch, prepare_snapshots): +def test__spcified_model_selector_not_existed( + epoch: int, prepare_snapshots: None +): with pytest.raises(FileNotFoundError): _ = select_snapshot_file( TEST_DATA_DIR, diff --git a/tests/test_nn/test_core_modules/test_concatenator.py b/tests/test_nn/test_core_modules/test_concatenator.py index e5c3981..84defc0 100644 --- a/tests/test_nn/test_core_modules/test_concatenator.py +++ b/tests/test_nn/test_core_modules/test_concatenator.py @@ -1,7 +1,6 @@ import numpy as np import pytest import torch - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import Concatenator @@ -21,7 +20,9 @@ def test__can_call_parameters(): ([(1, 2, 16), (1, 2, 16), (1, 2, 16)], (1, 2, 48)), ], ) -def test__concatenated_tensor_shape(input_shapes, desired_shape): +def test__concatenated_tensor_shape( + input_shapes: list[tuple[int]], desired_shape: tuple[int] +): phlower_tensors = { f"phlower_tensor_{i}": PhlowerTensor( torch.from_numpy(np.random.rand(*s)) diff --git a/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py index abdb50e..3e186ee 100644 --- a/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py +++ b/tests/test_nn/test_core_modules/test_en_equivariant_mlp.py @@ -1,12 +1,11 @@ import numpy as np import pytest import torch -from scipy.stats import ortho_group - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import EnEquivariantMLP from phlower.nn._core_modules import _functions +from scipy.stats import ortho_group def test__can_call_parameters(): @@ -34,7 +33,12 @@ def test__can_call_parameters(): ], ) @pytest.mark.parametrize("n_output_feature", [1, 16, 32]) -def test__en_equivariance(size, is_time_series, is_voxel, n_output_feature): +def test__en_equivariance( + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + n_output_feature: int, +): orthogonal_tensor = PhlowerTensor( torch.tensor(ortho_group.rvs(3).astype(np.float32)) ) diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index 3f4c6fd..a0c1a9c 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -1,12 +1,11 @@ import numpy as np import pytest import torch -from scipy import sparse as sp -from scipy.stats import ortho_group - from phlower import PhlowerTensor, phlower_tensor from phlower.nn._core_modules import _functions from phlower.utils.exceptions import PhlowerIncompatibleTensorError +from scipy import sparse as sp +from scipy.stats import ortho_group @pytest.mark.parametrize( @@ -26,7 +25,7 @@ ((4, 10, 3, 16), True, 5), ], ) -def test__spmm(size, is_time_series, repeat): +def test__spmm(size: tuple[int], is_time_series: bool, repeat: bool): phlower_tensor = PhlowerTensor( torch.rand(*size), is_time_series=is_time_series ) @@ -41,7 +40,7 @@ def test__spmm(size, is_time_series, repeat): sp_sparse = sp.coo_array(sparse.to_tensor().to_dense().numpy()) np_dense = phlower_tensor.to_tensor().numpy() - def assert_correct(actual, array): + def assert_correct(actual: np.ndarray, array: np.ndarray): dim_feat = len(array.shape) - 1 if dim_feat == 1: desired = array @@ -114,7 +113,11 @@ def test_smooth_leaky_relu_inverse(): ], ) def test_contraction_one_argument( - size, is_time_series, is_voxel, desired_pattern, dimension + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + desired_pattern: str, + dimension: list[list[int]], ): torch_tensor = torch.rand(*size) x = phlower_tensor( @@ -396,15 +399,15 @@ def test_contraction_one_argument( ], ) def test_contraction_two_arguments( - size_x, - size_y, - x_is_time_series, - y_is_time_series, - is_voxel, - desired_pattern, - desired_rank, - dimension_x, - dimension_y, + size_x: tuple[int], + size_y: tuple[int], + x_is_time_series: bool, + y_is_time_series: bool, + is_voxel: bool, + desired_pattern: str, + desired_rank: int, + dimension_x: list[list[int]] | None, + dimension_y: list[list[int]] | None, ): t_x = torch.rand(*size_x) x = phlower_tensor( @@ -665,14 +668,14 @@ def test_contraction_raises_phlower_incompatible_tensor_error(): ], ) def test_tensor_product( - size_x, - size_y, - x_is_time_series, - y_is_time_series, - is_voxel, - desired_pattern, - dimension_x, - dimension_y, + size_x: tuple[int], + size_y: tuple[int], + x_is_time_series: bool, + y_is_time_series: bool, + is_voxel: bool, + desired_pattern: str, + dimension_x: list[list[int]] | None, + dimension_y: list[list[int]] | None, ): t_x = torch.rand(*size_x) x = phlower_tensor( @@ -737,7 +740,11 @@ def test_tensor_product( ], ) def test_apply_orthogonal_group( - size, is_time_series, is_voxel, desired_pattern, dimension + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + desired_pattern: str, + dimension: list[list[int]] | None, ): orthogonal_matrix = torch.from_numpy(ortho_group.rvs(3).astype(np.float32)) orthogonal_tensor = phlower_tensor(orthogonal_matrix) @@ -800,7 +807,13 @@ def test_apply_orthogonal_group( [[-1], [-1], [2], [0], [1], [0], [0]], ], ) -def test_spatial_mean(size, is_time_series, is_voxel, mean_dims, dimension): +def test_spatial_mean( + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + mean_dims: int, + dimension: list[list[int]] | None, +): torch_tensor = torch.rand(*size) x = phlower_tensor( torch_tensor, diff --git a/tests/test_nn/test_core_modules/test_gcn.py b/tests/test_nn/test_core_modules/test_gcn.py index db767e4..e0855d3 100644 --- a/tests/test_nn/test_core_modules/test_gcn.py +++ b/tests/test_nn/test_core_modules/test_gcn.py @@ -1,6 +1,5 @@ import pytest import torch - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import GCN @@ -24,7 +23,7 @@ def test__can_call_parameters(): ((4, 10, 3, 16), True), ], ) -def test__gcn(size, is_time_series): +def test__gcn(size: tuple[int], is_time_series: bool): phlower_tensor = PhlowerTensor( torch.rand(*size), is_time_series=is_time_series ) diff --git a/tests/test_nn/test_core_modules/test_identity.py b/tests/test_nn/test_core_modules/test_identity.py index 5e07b1e..e3c945f 100644 --- a/tests/test_nn/test_core_modules/test_identity.py +++ b/tests/test_nn/test_core_modules/test_identity.py @@ -1,7 +1,6 @@ import numpy as np import pytest import torch - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import Identity @@ -21,7 +20,7 @@ def test__can_call_parameters(): (1, 2, 48), ], ) -def test__identity(input_shape): +def test__identity(input_shape: tuple[int]): phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) phlower_tensors = phlower_tensor_collection( {"phlower_tensor": phlower_tensor} diff --git a/tests/test_nn/test_core_modules/test_pinv.py b/tests/test_nn/test_core_modules/test_pinv.py index a9cabe5..fd70778 100644 --- a/tests/test_nn/test_core_modules/test_pinv.py +++ b/tests/test_nn/test_core_modules/test_pinv.py @@ -1,7 +1,6 @@ import numpy as np import pytest import torch - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import MLP, PInvMLP @@ -25,7 +24,7 @@ def test__can_call_parameters(): ([20, 20, 40, 100], ["tanh", "smooth_leaky_relu", "leaky_relu0p5"], 2), ], ) -def test__pinv_mlp(mlp_nodes, activations, decimal): +def test__pinv_mlp(mlp_nodes: list[int], activations: list[str], decimal: int): MLP0 = MLP(nodes=mlp_nodes, activations=activations) model = PInvMLP(reference_name="MLP0") diff --git a/tests/test_nn/test_core_modules/test_proportional.py b/tests/test_nn/test_core_modules/test_proportional.py index 3b51f05..919a694 100644 --- a/tests/test_nn/test_core_modules/test_proportional.py +++ b/tests/test_nn/test_core_modules/test_proportional.py @@ -1,7 +1,6 @@ import numpy as np import pytest import torch - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import Proportional @@ -27,7 +26,9 @@ def test__can_call_parameters(): ) @pytest.mark.parametrize("n_output_feature", [1, 16, 32]) @pytest.mark.parametrize("scale", [0.0, 0.5, 2.0]) -def test__proportional_linearity(size, is_time_series, n_output_feature, scale): +def test__proportional_linearity( + size: tuple[int], is_time_series: bool, n_output_feature: int, scale: float +): model = Proportional(nodes=[size[-1], n_output_feature]) phlower_tensor = PhlowerTensor( diff --git a/tests/test_nn/test_core_modules/test_share.py b/tests/test_nn/test_core_modules/test_share.py index 20c6d26..dc815a8 100644 --- a/tests/test_nn/test_core_modules/test_share.py +++ b/tests/test_nn/test_core_modules/test_share.py @@ -3,7 +3,6 @@ import numpy as np import pytest import torch - from phlower import PhlowerTensor from phlower.collections import phlower_tensor_collection from phlower.nn import MLP, Share @@ -23,7 +22,7 @@ def test__can_call_parameters(): "mlp_nodes", [([10, 10]), ([20, 10, 100])], ) -def test__reference_same_object(mlp_nodes): +def test__reference_same_object(mlp_nodes: list[int]): model = Share(reference_name="MLP0") MLP0 = MLP(nodes=mlp_nodes) model._reference = MLP0 @@ -41,7 +40,7 @@ def test__reference_same_object(mlp_nodes): @pytest.mark.parametrize("reference_name", ["MLP0", "GCP0"]) -def test__search_reference_name(reference_name): +def test__search_reference_name(reference_name: str): model = Share(reference_name=reference_name) mocked = mock.MagicMock(IReadonlyReferenceGroup) diff --git a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py index 5cb9b59..f607171 100644 --- a/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py +++ b/tests/test_nn/test_core_modules/test_similarity_equivariant_mlp.py @@ -1,8 +1,6 @@ import numpy as np import pytest import torch -from scipy.stats import ortho_group - from phlower import PhlowerTensor, phlower_tensor from phlower._base._dimension import PhysicalDimensions from phlower.collections import phlower_tensor_collection @@ -12,6 +10,7 @@ PhlowerDimensionRequiredError, PhlowerInvalidArgumentsError, ) +from scipy.stats import ortho_group def test__can_call_parameters(): @@ -56,14 +55,14 @@ def test__can_call_parameters(): @pytest.mark.parametrize("norm_function_name", ["identity", "sqrt"]) @pytest.mark.parametrize("centering", [False, True]) def test__similarity_equivariance( - size, - is_time_series, - is_voxel, - activation, - n_output_feature, - dimension, - norm_function_name, - centering, + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + activation: str, + n_output_feature: int, + dimension: dict, + norm_function_name: str, + centering: bool, ): orthogonal_tensor = PhlowerTensor( torch.tensor(ortho_group.rvs(3).astype(np.float32)) @@ -165,14 +164,14 @@ def test__similarity_equivariance( @pytest.mark.parametrize("norm_function_name", ["identity", "sqrt"]) @pytest.mark.parametrize("centering", [False, True]) def test__similarity_invariance( - size, - is_time_series, - is_voxel, - activation, - n_output_feature, - dimension, - norm_function_name, - centering, + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + activation: str, + n_output_feature: int, + dimension: dict, + norm_function_name: str, + centering: bool, ): orthogonal_tensor = PhlowerTensor( torch.tensor(ortho_group.rvs(3).astype(np.float32)) @@ -273,15 +272,15 @@ def test__similarity_invariance( @pytest.mark.parametrize("centering", [False, True]) @pytest.mark.parametrize("invariant", [False, True]) def test__scaling_equivariance( - size, - is_time_series, - is_voxel, - activation, - n_output_feature, - dimension, - norm_function_name, - centering, - invariant, + size: tuple[int], + is_time_series: bool, + is_voxel: bool, + activation: str, + n_output_feature: int, + dimension: dict, + norm_function_name: str, + centering: bool, + invariant: bool, ): orthogonal_tensor = PhlowerTensor(torch.eye(3)) dict_scaling_factor = { diff --git a/tests/test_nn/test_group_module.py b/tests/test_nn/test_group_module.py index 5a39c92..78b66bb 100644 --- a/tests/test_nn/test_group_module.py +++ b/tests/test_nn/test_group_module.py @@ -3,7 +3,6 @@ import numpy as np import pytest import scipy.sparse as sp - from phlower._base import phlower_array from phlower.collections import phlower_tensor_collection from phlower.nn import PhlowerGroupModule @@ -13,7 +12,7 @@ @pytest.mark.parametrize("yaml_file", ["with_share_nn.yml"]) -def test__resolve_modules_from_setting(yaml_file): +def test__resolve_modules_from_setting(yaml_file: str): setting_file = _SAMPLE_SETTING_DIR / yaml_file setting = PhlowerSetting.read_yaml(setting_file) @@ -22,7 +21,7 @@ def test__resolve_modules_from_setting(yaml_file): @pytest.mark.parametrize("yaml_file", ["with_share_nn.yml"]) -def test__draw(yaml_file): +def test__draw(yaml_file: str): output_directory = pathlib.Path(__file__).parent / "out" setting_file = _SAMPLE_SETTING_DIR / yaml_file @@ -37,7 +36,9 @@ def test__draw(yaml_file): @pytest.mark.parametrize( "yaml_file, input_n_feature, n_nodes", [("with_share_nn.yml", 10, 20)] ) -def test__forward_and_backward(yaml_file, input_n_feature, n_nodes): +def test__forward_and_backward( + yaml_file: str, input_n_feature: int, n_nodes: int +): setting_file = _SAMPLE_SETTING_DIR / yaml_file setting = PhlowerSetting.read_yaml(setting_file) diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_identity_scale.py b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_identity_scale.py index ffeb6f7..a648f21 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_identity_scale.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_identity_scale.py @@ -1,6 +1,5 @@ import numpy as np import pytest - from phlower.services.preprocessing._scalers.scale_functions import ( IdentityScaler, ) @@ -9,14 +8,14 @@ @pytest.mark.parametrize( "name, kwards", [("identity", {}), ("identity", {"some": 11})] ) -def test__create(name, kwards): +def test__create(name: str, kwards: dict): IdentityScaler.create(name, **kwards) @pytest.mark.parametrize( "name, kwards", [("isoam", {}), ("identities", {"some": 11})] ) -def test__cannot_create(name, kwards): +def test__cannot_create(name: str, kwards: dict): with pytest.raises(NotImplementedError): IdentityScaler.create(name, **kwards) @@ -29,7 +28,7 @@ def test__fixed_property_and_method(): @pytest.mark.parametrize("shape", [(3, 4), (4, 5), (2, 9)]) -def test__transform(shape): +def test__transform(shape: tuple[int]): scaler = IdentityScaler() data = np.random.rand(*shape) @@ -39,7 +38,7 @@ def test__transform(shape): @pytest.mark.parametrize("shape", [(3, 4), (4, 5), (2, 9)]) -def test__inverse_transform(shape): +def test__inverse_transform(shape: tuple[int]): scaler = IdentityScaler() data = np.random.rand(*shape) diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_isoam_scaler.py b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_isoam_scaler.py index 085bfa3..1bcabc4 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_isoam_scaler.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_isoam_scaler.py @@ -1,6 +1,5 @@ import numpy as np import pytest - from phlower.services.preprocessing._scalers.scale_functions import IsoAMScaler @@ -11,7 +10,7 @@ ("isoam_scale", {"other_components": ["a", "b"]}), ], ) -def test__create(name, kwards): +def test__create(name: str, kwards: dict): scaler = IsoAMScaler.create(name, **kwards) assert scaler.other_components == kwards["other_components"] @@ -23,7 +22,7 @@ def test__create(name, kwards): ("maxabs_scale", {"other_components": ["a", "b"]}), ], ) -def test__cannot_create(name, kwards): +def test__cannot_create(name: str, kwards: dict): with pytest.raises(NotImplementedError): IsoAMScaler.create(name, **kwards) @@ -36,7 +35,7 @@ def test__cannot_create(name, kwards): ([400, 40, 4], ["a", "b", "c"]), ], ) -def test__partial_fit(n_data_list, other_components): +def test__partial_fit(n_data_list: list[int], other_components: list[str]): scaler = IsoAMScaler(other_components=other_components) data_list = [np.random.rand(n_data) for n_data in n_data_list] @@ -51,7 +50,7 @@ def test__partial_fit(n_data_list, other_components): @pytest.mark.parametrize("n_data, std", [(100, 5.0), (300, 4.0)]) -def test__transform(n_data, std): +def test__transform(n_data: int, std: float): scaler = IsoAMScaler(other_components=["dummy"]) scaler.std_ = std @@ -69,7 +68,9 @@ def test__transform(n_data, std): ([400, 40, 4], ["a", "b", "c"]), ], ) -def test__retrieve_from_dumped_data(n_data_list, other_components): +def test__retrieve_from_dumped_data( + n_data_list: list[int], other_components: list[str] +): scaler = IsoAMScaler(other_components=other_components) data_list = [np.random.rand(n_data) for n_data in n_data_list] diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_max_abs_powered_scaler.py b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_max_abs_powered_scaler.py index 217f5df..4ce7b08 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_max_abs_powered_scaler.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_max_abs_powered_scaler.py @@ -1,19 +1,19 @@ import numpy as np import pytest import scipy.sparse as sp - from phlower.services.preprocessing._scalers.scale_functions import ( MaxAbsPoweredScaler, ) +from phlower.utils.typing import ArrayDataType @pytest.mark.parametrize("name, kwards", [("max_abs_powered", {})]) -def test__create(name, kwards): +def test__create(name: str, kwards: dict): _ = MaxAbsPoweredScaler.create(name, **kwards) @pytest.mark.parametrize("name", ["isoam_scale", "identity"]) -def test__failed_create(name): +def test__failed_create(name: str): with pytest.raises(NotImplementedError): MaxAbsPoweredScaler.create(name) @@ -63,7 +63,7 @@ def test__default_parameter(): ), ], ) -def test__partial_fit(data, desired): +def test__partial_fit(data: list[ArrayDataType], desired: list[float]): scaler = MaxAbsPoweredScaler.create("max_abs_powered") for v in data: scaler.partial_fit(v) @@ -107,7 +107,9 @@ def test__partial_fit(data, desired): ), ], ) -def test__transform(max_data, power, inputs, desired): +def test__transform( + max_data: np.ndarray, power: float, inputs: np.ndarray, desired: np.ndarray +): scaler = MaxAbsPoweredScaler.create("max_abs_powered", power=power) scaler.max_ = max_data @@ -149,7 +151,9 @@ def test__transform(max_data, power, inputs, desired): ), ], ) -def test__transform_sparse_array(max_data, power, inputs, desired): +def test__transform_sparse_array( + max_data: np.ndarray, power: float, inputs: np.ndarray, desired: np.ndarray +): scaler = MaxAbsPoweredScaler.create("max_abs_powered", power=power) scaler.max_ = max_data @@ -181,7 +185,9 @@ def test__transform_sparse_array(max_data, power, inputs, desired): ), ], ) -def test__inverse_transform(max_data, power, inputs): +def test__inverse_transform( + max_data: np.ndarray, power: float, inputs: np.ndarray +): scaler = MaxAbsPoweredScaler.create("max_abs_powered", power=power) scaler.max_ = max_data @@ -205,7 +211,9 @@ def test__inverse_transform(max_data, power, inputs): ), ], ) -def test__inverse_transform_sparse_array(max_data, power, inputs): +def test__inverse_transform_sparse_array( + max_data: np.ndarray, power: float, inputs: np.ndarray +): scaler = MaxAbsPoweredScaler.create("max_abs_powered", power=power) scaler.max_ = max_data @@ -221,7 +229,7 @@ def test__inverse_transform_sparse_array(max_data, power, inputs): ({"max_": [1.0, 3.0]}, sp.coo_matrix(np.array([[3.0], [7.0]]))), ], ) -def test__raise_error_when_transform(kwards, data): +def test__raise_error_when_transform(kwards: dict, data: ArrayDataType): scaler = MaxAbsPoweredScaler.create("max_abs_powered") for k, v in kwards.items(): setattr(scaler, k, np.array(v)) @@ -237,7 +245,9 @@ def test__raise_error_when_transform(kwards, data): ("max_abs_powered", 2.0, 100000, 3), ], ) -def test__retrieve_from_dumped_data(name, power, n_nodes, n_feature): +def test__retrieve_from_dumped_data( + name: str, power: float, n_nodes: int, n_feature: int +): interim_value = np.random.randn(n_nodes, n_feature) * 2 scaler = MaxAbsPoweredScaler.create(name, power=power) diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_min_max_scaler.py b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_min_max_scaler.py index 1d1f169..14f0ac8 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_min_max_scaler.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_min_max_scaler.py @@ -1,13 +1,12 @@ import numpy as np import pytest - from phlower.services.preprocessing._scalers.scale_functions import MinMaxScaler @pytest.mark.parametrize( "name, kwards", [("min_max", {}), ("min_max", {"copy": True})] ) -def test__create(name, kwards): +def test__create(name: str, kwards: dict): _ = MinMaxScaler.create(name, **kwards) @@ -15,7 +14,7 @@ def test__create(name, kwards): "name, kwards", [("identity", {}), ("isoam_scale", {"other_components": ["a"]})], ) -def test__invalid_name_create(name, kwards): +def test__invalid_name_create(name: str, kwards: dict): with pytest.raises(NotImplementedError): _ = MinMaxScaler.create(name, **kwards) @@ -33,7 +32,7 @@ def test__fixed_property_and_method(): ("min_max", 100000, 3), ], ) -def test__retrieve_from_dumped_data(name, n_nodes, n_feature): +def test__retrieve_from_dumped_data(name: str, n_nodes: int, n_feature: int): interim_value = np.random.randn(n_nodes, n_feature) * 2 scaler = MinMaxScaler.create(name) diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_standard_scaler.py b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_standard_scaler.py index 3dc5d41..59ed8ec 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_standard_scaler.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scale_functions/test_standard_scaler.py @@ -1,6 +1,5 @@ import numpy as np import pytest - from phlower.services.preprocessing._scalers.scale_functions import ( StandardScaler, ) @@ -9,7 +8,7 @@ @pytest.mark.parametrize( "name, kwards", [("std_scale", {}), ("standardize", {})] ) -def test__create(name, kwards): +def test__create(name: str, kwards: dict): _ = StandardScaler.create(name, **kwards) @@ -20,7 +19,7 @@ def test__create(name, kwards): ("standardize", {"with_mean": False}), ], ) -def test__invalid_args_create(name, kwards): +def test__invalid_args_create(name: str, kwards: dict): with pytest.raises(ValueError): _ = StandardScaler.create(name, **kwards) @@ -29,7 +28,7 @@ def test__invalid_args_create(name, kwards): "name, kwards", [("identity", {}), ("isoam_scale", {"other_components": ["a"]})], ) -def test__invalid_name_create(name, kwards): +def test__invalid_name_create(name: str, kwards: dict): with pytest.raises(NotImplementedError): _ = StandardScaler.create(name, **kwards) @@ -52,7 +51,7 @@ def test__fixed_property_and_method(): (100000, 5, [1, 2, 3, 4, 5]), ], ) -def test__std_scale_values(n_nodes, n_feature, means): +def test__std_scale_values(n_nodes: int, n_feature: int, means: list[int]): interim_value = np.random.randn(n_nodes, n_feature) * 2 + np.array([means]) scaler = StandardScaler.create(name="std_scale") @@ -76,7 +75,7 @@ def test__std_scale_values(n_nodes, n_feature, means): "n_nodes, n_feature, means", [(100000, 5, [1, 2, 3, 4, 5]), (100000, 3, [10, 21, 3])], ) -def test__standardize_values(n_nodes, n_feature, means): +def test__standardize_values(n_nodes: int, n_feature: int, means: list[int]): interim_value = np.random.randn(n_nodes, n_feature) * 2 + np.array([means]) scaler = StandardScaler() @@ -103,7 +102,7 @@ def test__standardize_values(n_nodes, n_feature, means): ("standardize", 100000, 3), ], ) -def test__retrieve_from_dumped_data(name, n_nodes, n_feature): +def test__retrieve_from_dumped_data(name: str, n_nodes: int, n_feature: int): interim_value = np.random.randn(n_nodes, n_feature) * 2 scaler = StandardScaler.create(name) diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scaler_wrapper.py b/tests/test_services/test_preprocessing/test_scalers/test_scaler_wrapper.py index 044ee9a..e1238d3 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scaler_wrapper.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scaler_wrapper.py @@ -1,11 +1,11 @@ import numpy as np import pytest import scipy.sparse as sp - from phlower.services.preprocessing._scalers import ( PhlowerScalerWrapper, scale_functions, ) +from phlower.utils.typing import DenseArrayType, SparseArrayType @pytest.fixture @@ -42,7 +42,9 @@ def sample_sparse_data() -> sp.csc_matrix: ("max_abs_powered", scale_functions.MaxAbsPoweredScaler), ], ) -def test__initialized_class(scaler_name, cls): +def test__initialized_class( + scaler_name: str, cls: scale_functions.IPhlowerScaler +): kwards = {} if scaler_name == "isoam_scale": kwards = {"other_components": ["a", "b"]} @@ -51,7 +53,9 @@ def test__initialized_class(scaler_name, cls): @pytest.mark.parametrize("scaler_name", ["std_scale", "max_abs_powered"]) -def test__transform_sparse_data(scaler_name, sample_sparse_data): +def test__transform_sparse_data( + scaler_name: str, sample_sparse_data: SparseArrayType +): scaler = PhlowerScalerWrapper(scaler_name) scaler.partial_fit(sample_sparse_data) transformed_data = scaler.transform(sample_sparse_data) @@ -61,7 +65,7 @@ def test__transform_sparse_data(scaler_name, sample_sparse_data): @pytest.mark.parametrize("scaler_name", ["std_scale", "standardize", "min_max"]) -def test__get_dumped_data(scaler_name, sample_data): +def test__get_dumped_data(scaler_name: str, sample_data: DenseArrayType): scaler = PhlowerScalerWrapper(scaler_name) scaler.partial_fit(sample_data) diff --git a/tests/test_services/test_preprocessing/test_scalers/test_scalers_composition.py b/tests/test_services/test_preprocessing/test_scalers/test_scalers_composition.py index 7c90ea4..cbf5578 100644 --- a/tests/test_services/test_preprocessing/test_scalers/test_scalers_composition.py +++ b/tests/test_services/test_preprocessing/test_scalers/test_scalers_composition.py @@ -3,7 +3,6 @@ import numpy as np import pytest - from phlower.io import PhlowerNumpyFile from phlower.services.preprocessing._scalers import ( PhlowerScalerWrapper, @@ -27,7 +26,7 @@ def test__from_setting(): ("scaler_c", True, None), ], ) -def test__get_scaler(name, allow_missing, method_name): +def test__get_scaler(name: str, allow_missing: bool, method_name: str): name2scaler = {"scaler_a": "identity", "scaler_b": "std_scale"} scalers_dict = {k: PhlowerScalerWrapper(v) for k, v in name2scaler.items()} @@ -43,7 +42,7 @@ def test__get_scaler(name, allow_missing, method_name): @pytest.fixture(scope="module") -def create_sample_dataset(): +def create_sample_dataset() -> list[PhlowerNumpyFile]: output_directory = pathlib.Path(__file__).parent / "tmp" if output_directory.exists(): @@ -76,7 +75,9 @@ def create_sample_dataset(): ) ], ) -def test__retrieve_from_dumped_data(name2scaler, create_sample_dataset): +def test__retrieve_from_dumped_data( + name2scaler: dict, create_sample_dataset: list[PhlowerNumpyFile] +): data_files = create_sample_dataset # fit to sample data diff --git a/tests/test_services/test_trainer/test_logging_items/test_logging_items.py b/tests/test_services/test_trainer/test_logging_items/test_logging_items.py index c621502..c614881 100644 --- a/tests/test_services/test_trainer/test_logging_items/test_logging_items.py +++ b/tests/test_services/test_trainer/test_logging_items/test_logging_items.py @@ -1,6 +1,6 @@ import pytest - from phlower.services.trainer.logging_items import ( + ILoggingItem, LoggingDictItem, LoggingFloatItem, LoggingIntItem, @@ -12,7 +12,7 @@ @pytest.mark.parametrize( "val, title, expected", [(3, "sample", "3"), (100, "sample", "100")] ) -def test__logging_int_item_to_str(val, title, expected): +def test__logging_int_item_to_str(val: int, title: str, expected: str): item = LoggingIntItem(val=val, title=title) assert item.format() == expected @@ -28,7 +28,7 @@ def test__logging_int_item_to_str(val, title, expected): ], ) def test__logging_int_item_to_format_str( - val, title, margin, force_title, expected + val: int, title: str, margin: int, force_title: str | None, expected: str ): item = LoggingIntItem(val=val, title=title) actual = item.format(padding_margin=margin, title=force_title) @@ -43,7 +43,7 @@ def test__logging_int_item_to_format_str( (0.0000123456, "sample", ".5e"), ], ) -def test__logging_float_item_to_str(val, title, formatter): +def test__logging_float_item_to_str(val: float, title: str, formatter: str): item = LoggingFloatItem(val=val, title=title) actual = item.format(formatter=formatter) expected = f"{val:{formatter}}" @@ -61,7 +61,7 @@ def test__logging_float_item_to_str(val, title, formatter): ], ) def test__logging_float_item_to_format_str( - val, title, formatter, force_title, margin + val: float, title: str, formatter: str, force_title: str | None, margin: int ): item = LoggingFloatItem(val=val, title=title) actual = item.format( @@ -83,7 +83,7 @@ def test__logging_float_item_to_format_str( ("sample_test", "sample_test"), ], ) -def test__logging_str_item_to_str(val, expected): +def test__logging_str_item_to_str(val: str, expected: str): item = LoggingStrItem(val=val) actual = item.format() assert actual == expected @@ -97,7 +97,7 @@ def test__logging_str_item_to_str(val, expected): ("sample_test", 4), ], ) -def test__logging_str_item_to_format_str(val, margin): +def test__logging_str_item_to_format_str(val: str, margin: int): item = LoggingStrItem(val=val) actual = item.format(padding_margin=margin) expected = val + " " * margin @@ -111,7 +111,9 @@ def test__logging_str_item_to_format_str(val, margin): ({"v1": 1234455, "v2": 7832989}, "samples/", ".2e"), ], ) -def test__logging_dict_item_to_str(val, title, formatter): +def test__logging_dict_item_to_str( + val: dict[str, float], title: str, formatter: str +): item = LoggingDictItem(val=val, title=title) actual = item.format(formatter=formatter) @@ -130,7 +132,11 @@ def test__logging_dict_item_to_str(val, title, formatter): ], ) def test__logging_dict_item_to_format_str( - val, title, formatter, margin, force_title + val: dict[str, float], + title: str, + formatter: str, + margin: int, + force_title: str | None, ): item = LoggingDictItem(val=val, title=title) actual = item.format( @@ -154,7 +160,9 @@ def test__logging_dict_item_to_format_str( ({}, "samples/", ".2e", 8), ], ) -def test__logging_empty_dict_item_to_format_str(val, title, formatter, margin): +def test__logging_empty_dict_item_to_format_str( + val: dict, title: str, formatter: str, margin: int +): item = LoggingDictItem(val=val, title=title) actual = item.format(formatter=formatter, padding_margin=margin) assert actual == "" @@ -169,7 +177,9 @@ def test__logging_empty_dict_item_to_format_str(val, title, formatter, margin): ({"a": 3.14}, "test", LoggingDictItem), ], ) -def test__create_logitems(val, title, expected): +def test__create_logitems( + val: int | float | dict, title: str, expected: ILoggingItem +): item = create_logitems(value=val, title=title) assert isinstance(item, expected) diff --git a/tests/test_services/test_trainer/test_optimizer.py b/tests/test_services/test_trainer/test_optimizer.py index a18122e..362b8cc 100644 --- a/tests/test_services/test_trainer/test_optimizer.py +++ b/tests/test_services/test_trainer/test_optimizer.py @@ -2,7 +2,6 @@ import pytest import torch - from phlower.services.trainer._optimizer import PhlowerOptimizerWrapper from phlower.settings import PhlowerTrainerSetting @@ -25,7 +24,7 @@ ], ) def test__pass_kwargs_when_call_from_setting( - optimizer, optimizer_parameters, schedulers + optimizer: str, optimizer_parameters: dict, schedulers: list[dict] ): setting = PhlowerTrainerSetting( loss_setting={"name2loss": {}}, @@ -60,7 +59,12 @@ def test__pass_kwargs_when_call_from_setting( ("SGD", 0.0005, 0.01, torch.optim.SGD), ], ) -def test__optimizer_parameters(optimizer, lr, weight_decay, desired_optimizer): +def test__optimizer_parameters( + optimizer: str, + lr: float, + weight_decay: float, + desired_optimizer: type[torch.optim.Optimizer], +): model = torch.nn.Linear(in_features=10, out_features=10) optimizer = PhlowerOptimizerWrapper( parameters=model.parameters(), @@ -102,7 +106,9 @@ def test__optimizer_parameters(optimizer, lr, weight_decay, desired_optimizer): ), ], ) -def test__scheduler_parameters(schedulers, desired): +def test__scheduler_parameters( + schedulers: dict, desired: list[type[torch.optim.Optimizer]] +): dummy = torch.nn.Linear(in_features=10, out_features=10) optimizer = PhlowerOptimizerWrapper( parameters=dummy.parameters(), diff --git a/tests/test_settings/test_model_settings.py b/tests/test_settings/test_model_settings.py index ce88a12..7fce089 100644 --- a/tests/test_settings/test_model_settings.py +++ b/tests/test_settings/test_model_settings.py @@ -2,7 +2,6 @@ import pytest import yaml - from phlower.settings import GroupModuleSetting, ModuleSetting from phlower.utils.exceptions import ( PhlowerModuleCycleError, @@ -39,7 +38,7 @@ def _recursive_check( @pytest.mark.parametrize( "file_name", ["simple_module.yml", "simple_group_in_group.yml"] ) -def test__can_resolve_phlower_networks(file_name): +def test__can_resolve_phlower_networks(file_name: str): data = parse_file(file_name) setting = GroupModuleSetting(**data["model"]) @@ -49,7 +48,7 @@ def test__can_resolve_phlower_networks(file_name): @pytest.mark.parametrize("file_name", ["cycle_error.yml"]) -def test__detect_cycle_error(file_name): +def test__detect_cycle_error(file_name: str): data = parse_file(file_name) setting = GroupModuleSetting(**data["model"]) @@ -58,7 +57,7 @@ def test__detect_cycle_error(file_name): @pytest.mark.parametrize("file_name", ["not_matched_last_node_error.yml"]) -def test__detect_ndim_inconsistency(file_name): +def test__detect_ndim_inconsistency(file_name: str): data = parse_file(file_name) setting = GroupModuleSetting(**data["model"]) @@ -67,7 +66,7 @@ def test__detect_ndim_inconsistency(file_name): @pytest.mark.parametrize("file_name", ["duplicate_keys_error.yml"]) -def test__detect_duplicate_errors(file_name): +def test__detect_duplicate_errors(file_name: str): data = parse_file(file_name) setting = GroupModuleSetting(**data["model"]) @@ -76,7 +75,7 @@ def test__detect_duplicate_errors(file_name): @pytest.mark.parametrize("file_name", ["key_missing_error.yml"]) -def test__detect_key_missing(file_name): +def test__detect_key_missing(file_name: str): data = parse_file(file_name) setting = GroupModuleSetting(**data["model"]) diff --git a/tests/test_settings/test_module_settings/test_concatenator_setting.py b/tests/test_settings/test_module_settings/test_concatenator_setting.py index 5034252..d62f73d 100644 --- a/tests/test_settings/test_module_settings/test_concatenator_setting.py +++ b/tests/test_settings/test_module_settings/test_concatenator_setting.py @@ -1,21 +1,21 @@ import pathlib +from collections.abc import Callable import hypothesis.strategies as st import pytest import yaml from hypothesis import assume, given, settings - from phlower.settings import PhlowerModelSetting from phlower.settings._module_settings import ConcatenatorSetting @pytest.mark.parametrize("nodes", [(None), ([10, 10])]) -def test__can_accept_valid_n_nodes(nodes): +def test__can_accept_valid_n_nodes(nodes: list[int] | None): _ = ConcatenatorSetting(nodes=nodes) @pytest.mark.parametrize("nodes", [([5]), ([10, 10, 10])]) -def test__raise_error_when_invalid_n_nodes(nodes): +def test__raise_error_when_invalid_n_nodes(nodes: list[int]): with pytest.raises(ValueError): _ = ConcatenatorSetting(nodes=nodes) @@ -23,14 +23,14 @@ def test__raise_error_when_invalid_n_nodes(nodes): @pytest.mark.parametrize( "input_dims, desired", [([30, 50, 40], 120), ([40], 40), ([100, 10], 110)] ) -def test__gather_input_dims(input_dims, desired): +def test__gather_input_dims(input_dims: list[int], desired: int): setting = ConcatenatorSetting() assert setting.gather_input_dims(*input_dims) == desired @st.composite -def same_length_lists(draw): +def same_length_lists(draw: Callable) -> tuple[list[int]]: n_elements = draw(st.integers(min_value=2, max_value=2)) fixed_length_list = st.lists( st.integers(min_value=1, max_value=200), @@ -43,7 +43,9 @@ def same_length_lists(draw): @given(same_length_lists()) @settings(max_examples=100) -def test__nodes_is_update_after_overwrite_nodes(lists): +def test__nodes_is_update_after_overwrite_nodes( + lists: tuple[list[int], list[int]], +): nodes, update_nodes = lists assume(nodes != update_nodes) setting = ConcatenatorSetting() @@ -67,7 +69,7 @@ def test__reference_is_not_necessary(): @pytest.mark.parametrize("yaml_file", ["check_concatenator_nodes.yml"]) -def test__nodes_after_resolve(yaml_file): +def test__nodes_after_resolve(yaml_file: str): with open(_TEST_DATA_DIR / yaml_file) as fr: content = yaml.load(fr, Loader=yaml.SafeLoader) diff --git a/tests/test_settings/test_module_settings/test_gcn_settings.py b/tests/test_settings/test_module_settings/test_gcn_settings.py index 0061f57..24b7d37 100644 --- a/tests/test_settings/test_module_settings/test_gcn_settings.py +++ b/tests/test_settings/test_module_settings/test_gcn_settings.py @@ -1,10 +1,10 @@ import pathlib +from collections.abc import Callable import hypothesis.strategies as st import pytest import yaml from hypothesis import assume, given, settings - from phlower.settings import PhlowerModelSetting from phlower.settings._module_settings import GCNSetting @@ -18,7 +18,9 @@ ([-1, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), ], ) -def test__can_accept_valid_n_nodes(nodes, activations, dropouts): +def test__can_accept_valid_n_nodes( + nodes: list[int], activations: list[str], dropouts: list[float] +): _ = GCNSetting( nodes=nodes, support_name="dummy", @@ -38,7 +40,9 @@ def test__can_accept_valid_n_nodes(nodes, activations, dropouts): ([10], [], []), ], ) -def test__raise_error_when_invalid_n_nodes(nodes, activations, dropouts): +def test__raise_error_when_invalid_n_nodes( + nodes: list[int], activations: list[str], dropouts: list[float] +): with pytest.raises(ValueError): _ = GCNSetting( nodes=nodes, @@ -49,7 +53,7 @@ def test__raise_error_when_invalid_n_nodes(nodes, activations, dropouts): @pytest.mark.parametrize("input_dims", [([30]), ([40]), ([100])]) -def test__gather_input_dims(input_dims): +def test__gather_input_dims(input_dims: list[int]): setting = GCNSetting( nodes=[10, 20], support_name="dummy", @@ -61,7 +65,7 @@ def test__gather_input_dims(input_dims): @pytest.mark.parametrize("input_dims", [([]), ([40, 400]), ([10, 0, 1])]) -def test__raise_error_invalid_input_dims(input_dims): +def test__raise_error_invalid_input_dims(input_dims: list[int]): setting = GCNSetting( nodes=[10, 20], support_name="dummy", @@ -77,7 +81,9 @@ def test__raise_error_invalid_input_dims(input_dims): "nodes, activations, dropouts", [([10, 20, 30], [], []), ([10, 20, 30, 40, 50], [], [])], ) -def test__fill_default_settings(nodes, activations, dropouts): +def test__fill_default_settings( + nodes: list[int], activations: list[str], dropouts: list[float] +): setting = GCNSetting( nodes=nodes, support_name="dummy", @@ -92,7 +98,7 @@ def test__fill_default_settings(nodes, activations, dropouts): @st.composite -def same_length_lists(draw): +def same_length_lists(draw: Callable) -> tuple[list[int], list[int]]: n_elements = draw(st.integers(min_value=2, max_value=10)) fixed_length_list = st.lists( st.integers(min_value=1, max_value=200), @@ -105,7 +111,9 @@ def same_length_lists(draw): @given(same_length_lists()) @settings(max_examples=100) -def test__nodes_is_update_after_overwrite_nodes(lists): +def test__nodes_is_update_after_overwrite_nodes( + lists: tuple[list[int], list[int]], +): nodes, update_nodes = lists assume(nodes != update_nodes) setting = GCNSetting( @@ -139,7 +147,7 @@ def test__reference_is_not_necessary(): @pytest.mark.parametrize("yaml_file", ["check_gcn_nodes.yml"]) -def test__nodes_after_resolve(yaml_file): +def test__nodes_after_resolve(yaml_file: str): with open(_TEST_DATA_DIR / yaml_file) as fr: content = yaml.load(fr, Loader=yaml.SafeLoader) diff --git a/tests/test_settings/test_module_settings/test_mlp_setting.py b/tests/test_settings/test_module_settings/test_mlp_setting.py index 0835a2f..29d3f9e 100644 --- a/tests/test_settings/test_module_settings/test_mlp_setting.py +++ b/tests/test_settings/test_module_settings/test_mlp_setting.py @@ -1,10 +1,10 @@ import pathlib +from collections.abc import Callable import hypothesis.strategies as st import pytest import yaml from hypothesis import assume, given, settings - from phlower.settings import PhlowerModelSetting from phlower.settings._module_settings import MLPSetting @@ -20,7 +20,9 @@ ([-1, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), ], ) -def test__can_accept_valid_n_nodes(nodes, activations, dropouts): +def test__can_accept_valid_n_nodes( + nodes: list[int], activations: list[str], dropouts: list[float] +): _ = MLPSetting(nodes=nodes, activations=activations, dropouts=dropouts) @@ -35,7 +37,9 @@ def test__can_accept_valid_n_nodes(nodes, activations, dropouts): ([10], [], []), ], ) -def test__raise_error_when_invalid_n_nodes(nodes, activations, dropouts): +def test__raise_error_when_invalid_n_nodes( + nodes: list[int], activations: list[str], dropouts: list[float] +): with pytest.raises(ValueError): _ = MLPSetting(nodes=nodes, activations=activations, dropouts=dropouts) @@ -44,7 +48,9 @@ def test__raise_error_when_invalid_n_nodes(nodes, activations, dropouts): "nodes, activations, dropouts", [([10, 20, 30], [], []), ([10, 20, 30, 40, 50], [], [])], ) -def test__fill_default_settings(nodes, activations, dropouts): +def test__fill_default_settings( + nodes: list[int], activations: list[str], dropouts: list[float] +): setting = MLPSetting( nodes=nodes, activations=activations, dropouts=dropouts ) @@ -56,7 +62,7 @@ def test__fill_default_settings(nodes, activations, dropouts): @pytest.mark.parametrize("input_dims", [([30]), ([40]), ([100])]) -def test__gather_input_dims(input_dims): +def test__gather_input_dims(input_dims: list[int]): setting = MLPSetting( nodes=[10, 20], activations=["identity"], dropouts=[0.1] ) @@ -65,7 +71,7 @@ def test__gather_input_dims(input_dims): @pytest.mark.parametrize("input_dims", [([]), ([40, 400]), ([10, 0, 1])]) -def test__raise_error_invalid_input_dims(input_dims): +def test__raise_error_invalid_input_dims(input_dims: list[int]): setting = MLPSetting( nodes=[10, 20], activations=["identity"], dropouts=[0.1] ) @@ -75,7 +81,7 @@ def test__raise_error_invalid_input_dims(input_dims): @st.composite -def same_length_lists(draw): +def same_length_lists(draw: Callable) -> tuple[list[int], list[int]]: n_elements = draw(st.integers(min_value=2, max_value=10)) fixed_length_list = st.lists( st.integers(min_value=1, max_value=200), @@ -88,7 +94,9 @@ def same_length_lists(draw): @given(same_length_lists()) @settings(max_examples=100) -def test__nodes_is_update_after_overwrite_nodes(lists): +def test__nodes_is_update_after_overwrite_nodes( + lists: tuple[list[int], list[int]], +): nodes, update_nodes = lists assume(nodes != update_nodes) setting = MLPSetting( @@ -118,7 +126,7 @@ def test__reference_is_not_necessary(): @pytest.mark.parametrize("yaml_file", ["check_mlp_nodes.yml"]) -def test__nodes_after_resolve(yaml_file): +def test__nodes_after_resolve(yaml_file: str): with open(_TEST_DATA_DIR / yaml_file) as fr: content = yaml.load(fr, Loader=yaml.SafeLoader) diff --git a/tests/test_settings/test_module_settings/test_share_setting.py b/tests/test_settings/test_module_settings/test_share_setting.py index d842efa..94b2689 100644 --- a/tests/test_settings/test_module_settings/test_share_setting.py +++ b/tests/test_settings/test_module_settings/test_share_setting.py @@ -5,14 +5,13 @@ import pytest import yaml from hypothesis import given, settings - from phlower.settings import PhlowerModelSetting, PhlowerSetting from phlower.settings._module_settings import ShareSetting @given(st.lists(st.integers(), max_size=100)) @settings(max_examples=100) -def test__gather_input_dims_of_reference_setting(input_dims): +def test__gather_input_dims_of_reference_setting(input_dims: list[int]): setting = ShareSetting(reference_name="dummy") mocked = mock.MagicMock() setting.reference = mocked @@ -33,7 +32,7 @@ def test__get_n_nodes_of_reference_setting(): @pytest.mark.parametrize("reference_name", ["dummy", "mlp0", "gcn0"]) -def test__call_parent_function_when_get_reference(reference_name): +def test__call_parent_function_when_get_reference(reference_name: str): setting = ShareSetting(reference_name=reference_name) mocked = mock.MagicMock() @@ -55,7 +54,7 @@ def test__can_resolve(): @pytest.mark.parametrize( "yaml_file", ["check_gcn_share_nodes.yml", "check_mlp_share_nodes.yml"] ) -def test__nodes_after_resolve(yaml_file): +def test__nodes_after_resolve(yaml_file: str): with open(_TEST_DATA_DIR / yaml_file) as fr: content = yaml.load(fr, Loader=yaml.SafeLoader) diff --git a/tests/test_settings/test_phlower_setting.py b/tests/test_settings/test_phlower_setting.py index 350424f..cd704de 100644 --- a/tests/test_settings/test_phlower_setting.py +++ b/tests/test_settings/test_phlower_setting.py @@ -3,7 +3,6 @@ import pydantic import pytest - from phlower.io import PhlowerYamlFile from phlower.settings import ( PhlowerModelSetting, @@ -16,7 +15,7 @@ "variable_dimensions", [({"vala": {"mass_": 1, "time": 4}}), ({"sample": {"kg": 2, "Am": -2}})], ) -def test__invalid_dimension_model_setting(variable_dimensions): +def test__invalid_dimension_model_setting(variable_dimensions: dict): with pytest.raises(TypeError): _ = PhlowerModelSetting( variable_dimensions=variable_dimensions, @@ -32,12 +31,12 @@ def test__invalid_dimension_model_setting(variable_dimensions): @pytest.mark.parametrize( "selection_mode", ["best", "latest", "train_best", "specified"] ) -def test__valid_selection_mode(selection_mode): +def test__valid_selection_mode(selection_mode: str): _ = PhlowerPredictorSetting(selection_mode=selection_mode) @pytest.mark.parametrize("selection_mode", ["other", "best_of_best"]) -def test__invalid_selection_mode(selection_mode): +def test__invalid_selection_mode(selection_mode: str): with pytest.raises(pydantic.ValidationError): _ = PhlowerPredictorSetting(selection_mode=selection_mode) diff --git a/tests/test_settings/test_scaling_setting.py b/tests/test_settings/test_scaling_setting.py index 843c61d..1278114 100644 --- a/tests/test_settings/test_scaling_setting.py +++ b/tests/test_settings/test_scaling_setting.py @@ -3,8 +3,6 @@ import pydantic import pytest -from pipe import where - from phlower.io import PhlowerDirectory from phlower.settings import PhlowerScalingSetting from phlower.settings._scaling_setting import ( @@ -12,6 +10,7 @@ ScalerInputParameters, ScalerResolvedParameter, ) +from pipe import where # region test for ScalerInputParameters @@ -33,7 +32,7 @@ def test__is_parent(): ("user_defined", "val_b", "SCALER_val_b"), ], ) -def test_scaler_name(method, variable_name, desired): +def test_scaler_name(method: str, variable_name: str, desired: str): scaler = ScalerInputParameters(method=method) assert desired == scaler.get_scaler_name(variable_name) @@ -49,7 +48,7 @@ def test__is_not_parent(): @pytest.mark.parametrize("join_fitting", [True, False, None]) -def test__join_fitting(join_fitting): +def test__join_fitting(join_fitting: bool): if join_fitting is None: scaler = SameAsInputParameters(same_as="val_a") assert not scaler.join_fitting @@ -62,7 +61,7 @@ def test__join_fitting(join_fitting): @pytest.mark.parametrize( "same_as, desired", [("val_a", "SCALER_val_a"), ("val_b", "SCALER_val_b")] ) -def test_sameas_scaler_name(same_as, desired): +def test_sameas_scaler_name(same_as: str, desired: str): scaler = SameAsInputParameters(same_as=same_as) # NOTE: scaler name is decided by parent variable name @@ -118,7 +117,9 @@ def test__read_yml(): ({}, []), ], ) -def test__get_variable_names(scalers, desired): +def test__get_variable_names( + scalers: dict[str, dict[str, str]], desired: list[str] +): scaler = PhlowerScalingSetting(varaible_name_to_scalers=scalers) actual = scaler.get_variable_names() @@ -146,7 +147,11 @@ def test__get_variable_names(scalers, desired): ), ], ) -def test__is_scaler_exist(scalers, existed, not_existed): +def test__is_scaler_exist( + scalers: dict[str, dict[str, str]], + existed: list[str], + not_existed: list[str], +): scaler = PhlowerScalingSetting(varaible_name_to_scalers=scalers) for name in existed: @@ -174,7 +179,9 @@ def test__is_scaler_exist(scalers, existed, not_existed): ), ], ) -def test__get_scaler_name(scalers, desired): +def test__get_scaler_name( + scalers: dict[str, dict[str, str]], desired: list[tuple[str]] +): scaler = PhlowerScalingSetting(varaible_name_to_scalers=scalers) for key, ans in desired: @@ -187,7 +194,7 @@ def test__get_scaler_name(scalers, desired): @pytest.fixture -def create_sample_setting(): +def create_sample_setting() -> PhlowerScalingSetting: scalers = { "nodal": {"method": "std_scale", "parameters": {"std_": 0.001}}, "nodal_child": {"same_as": "nodal"}, @@ -202,7 +209,7 @@ def create_sample_setting(): return PhlowerScalingSetting(varaible_name_to_scalers=scalers) -def test__n_resolved_settings(create_sample_setting): +def test__n_resolved_settings(create_sample_setting: PhlowerScalingSetting): scaler: PhlowerScalingSetting = create_sample_setting resolved = scaler.resolve_scalers() @@ -217,8 +224,12 @@ def test__n_resolved_settings(create_sample_setting): ("SCALER_value_z", ["value_z"]), ], ) -def test__transform_items(scaler_name, desired, create_sample_setting): - scaler: PhlowerScalingSetting = create_sample_setting +def test__transform_items( + scaler_name: str, + desired: list[str], + create_sample_setting: PhlowerScalingSetting, +): + scaler = create_sample_setting resolved = scaler.resolve_scalers() target = list(resolved | where(lambda x: x.scaler_name == scaler_name))[0] @@ -234,7 +245,9 @@ def test__transform_items(scaler_name, desired, create_sample_setting): ("SCALER_value_z", ["value_z"]), ], ) -def test__fitting_items(scaler_name, desired, create_sample_setting): +def test__fitting_items( + scaler_name: str, desired: list[str], create_sample_setting: str +): scaler: PhlowerScalingSetting = create_sample_setting resolved = scaler.resolve_scalers() @@ -251,7 +264,9 @@ def test__fitting_items(scaler_name, desired, create_sample_setting): ("SCALER_value_y", True), ], ) -def test__component_wise(scaler_name, desired, create_sample_setting): +def test__component_wise( + scaler_name: str, desired: str, create_sample_setting: PhlowerScalingSetting +): scaler: PhlowerScalingSetting = create_sample_setting resolved = scaler.resolve_scalers() @@ -267,7 +282,11 @@ def test__component_wise(scaler_name, desired, create_sample_setting): ("SCALER_value_z", {"user_std_": 10.0}), ], ) -def test__parameters(scaler_name, desired, create_sample_setting): +def test__parameters( + scaler_name: str, + desired: dict[str, float], + create_sample_setting: PhlowerScalingSetting, +): scaler: PhlowerScalingSetting = create_sample_setting resolved = scaler.resolve_scalers() @@ -281,7 +300,9 @@ def test__parameters(scaler_name, desired, create_sample_setting): [("SCALER_nodal", True), ("SCALER_value_z", False)], ) def test__collect_fitting_files( - scaler_name, allow_missing, create_sample_setting + scaler_name: str, + allow_missing: bool, + create_sample_setting: PhlowerScalingSetting, ): scaler: PhlowerScalingSetting = create_sample_setting resolved = scaler.resolve_scalers() @@ -309,7 +330,9 @@ def test__collect_fitting_files( [("SCALER_nodal", True), ("SCALER_value_z", False)], ) def test__collect_transform_files( - scaler_name, allow_missing, create_sample_setting + scaler_name: str, + allow_missing: bool, + create_sample_setting: PhlowerScalingSetting, ): scaler: PhlowerScalingSetting = create_sample_setting resolved = scaler.resolve_scalers() @@ -360,7 +383,7 @@ def test__collect_transform_files( ), ], ) -def test__validate_isoam(scalers): +def test__validate_isoam(scalers: dict): setting = PhlowerScalingSetting(varaible_name_to_scalers=scalers) with pytest.raises(ValueError): _ = setting.resolve_scalers()[0] diff --git a/tests/test_utils/test_encryption.py b/tests/test_utils/test_encryption.py index 88d9542..6e6ceba 100644 --- a/tests/test_utils/test_encryption.py +++ b/tests/test_utils/test_encryption.py @@ -4,7 +4,6 @@ import shutil import pytest - from phlower.utils import decrypt_file, encrypt_file _OUTPUT_DIR = pathlib.Path(__file__).parent / "_tmp/_encryption" @@ -19,7 +18,7 @@ def prepare_empty_direcotry(): _OUTPUT_DIR.mkdir(parents=True) -def test__encrypt_file(prepare_empty_direcotry): +def test__encrypt_file(prepare_empty_direcotry: None): file_path = _OUTPUT_DIR / "sample.txt" encrypt_file( TEST_ENCRYPT_KEY, @@ -32,7 +31,7 @@ def test__encrypt_file(prepare_empty_direcotry): _ = fr.read() -def test__decrypt_file(prepare_empty_direcotry): +def test__decrypt_file(prepare_empty_direcotry: None): file_path = _OUTPUT_DIR / "sample.txt" content = "sample_content" encrypt_file( diff --git a/tests/test_utils/test_env.py b/tests/test_utils/test_env.py index 96ef6f8..7c566c3 100644 --- a/tests/test_utils/test_env.py +++ b/tests/test_utils/test_env.py @@ -1,7 +1,6 @@ from unittest import mock import pytest - from phlower.utils import determine_n_process @@ -9,7 +8,9 @@ "n_cpu, max_process, desired", [(10, 3, 3), (1, 3, 1), (10, None, 10), (0, None, 0)], ) -def test__determine_n_process(n_cpu, max_process, desired): +def test__determine_n_process( + n_cpu: int, max_process: int | None, desired: int +): with mock.patch("os.cpu_count", return_value=n_cpu): n_process = determine_n_process(max_process) diff --git a/tests/test_utils/test_multiprocessor.py b/tests/test_utils/test_multiprocessor.py index 1f29b35..8124c0f 100644 --- a/tests/test_utils/test_multiprocessor.py +++ b/tests/test_utils/test_multiprocessor.py @@ -1,10 +1,10 @@ import os import sys import time +from collections.abc import Callable import numpy as np import pytest - from phlower.utils import PhlowerMultiprocessor from phlower.utils._multiprocessor import _get_chunks, _process_chunk from phlower.utils.exceptions import PhlowerMultiProcessError @@ -17,7 +17,9 @@ (lambda x, y: x * 10 + y, [(1, 1), (2, 2), (3, 5)], [11, 22, 35]), ], ) -def test__process_chunk(fn, chunk, expects): +def test__process_chunk( + fn: Callable, chunk: list[tuple[int]], expects: list[int] +): actual = _process_chunk(fn, chunk) assert actual == expects @@ -41,7 +43,7 @@ def test__process_chunk(fn, chunk, expects): ), ], ) -def test__get_chunks(iterables, chunksize, expects): +def test__get_chunks(iterables: list, chunksize: int, expects: list): for i, chunk in enumerate(_get_chunks(iterables, chunksize=chunksize)): assert chunk == expects[i] @@ -56,12 +58,12 @@ def test__get_chunks(iterables, chunksize, expects): ) ], ) -def test__get_chunks_multiples(iterables, chunksize, expects): +def test__get_chunks_multiples(iterables: list, chunksize: int, expects: list): for i, chunk in enumerate(_get_chunks(*iterables, chunksize=chunksize)): assert chunk == expects[i] -def freaky_job(num: int): +def freaky_job(num: int) -> int: if num == 1: sys.exit(0) else: @@ -69,7 +71,7 @@ def freaky_job(num: int): @pytest.mark.parametrize("inputs, expects", [([3, 5, 6], [3, 5, 6])]) -def test__can_execute_functions(inputs, expects): +def test__can_execute_functions(inputs: list[int], expects: list[int]): processor = PhlowerMultiprocessor(max_process=2) results = processor.run(inputs, target_fn=freaky_job) @@ -77,13 +79,13 @@ def test__can_execute_functions(inputs, expects): @pytest.mark.parametrize("inputs", [([3, 1, 6]), ([1, 1, 1])]) -def test__can_detect_child_process_error(inputs): +def test__can_detect_child_process_error(inputs: list[int]): with pytest.raises(PhlowerMultiProcessError): processor = PhlowerMultiprocessor(max_process=2) _ = processor.run(inputs, target_fn=freaky_job) -def sample_sleep_job(num: int): +def sample_sleep_job(num: int) -> int: time.sleep(num) return num @@ -95,7 +97,7 @@ def sample_sleep_job(num: int): (2, [2, 2, 2, 2], 4), ], ) -def test__can_use_multi_core(max_process, inputs, expects): +def test__can_use_multi_core(max_process: int, inputs: list[int], expects: int): cpu_count = os.cpu_count() assert cpu_count >= max_process @@ -115,7 +117,9 @@ def test__can_use_multi_core(max_process, inputs, expects): "max_process, inputs, chunksize, expects", [(2, [2, 2, 2, 2], 3, 6), (2, [1, 1, 1, 1], 2, 2)], ) -def test__can_consider_chunksize(max_process, inputs, chunksize, expects): +def test__can_consider_chunksize( + max_process: int, inputs: list, chunksize: int, expects: int +): cpu_count = os.cpu_count() assert cpu_count >= max_process @@ -131,7 +135,7 @@ def test__can_consider_chunksize(max_process, inputs, chunksize, expects): np.testing.assert_approx_equal(elapsed_time, expects, significant=1) -def sample_add(num1: int, num2: int): +def sample_add(num1: int, num2: int) -> int: return num1 + num2 @@ -143,7 +147,12 @@ def sample_add(num1: int, num2: int): (2, [[1, 1, 1, 1], [2, 3, 5, 6]], 2, [3, 4, 6, 7]), ], ) -def test__can_flatten_return_objects(max_process, inputs, chunksize, expects): +def test__can_flatten_return_objects( + max_process: int, + inputs: list[list[int]], + chunksize: int, + expects: list[int], +): cpu_count = os.cpu_count() assert cpu_count >= max_process diff --git a/tests/test_utils/test_optimizer.py b/tests/test_utils/test_optimizer.py index 6e6e95b..8fc4469 100644 --- a/tests/test_utils/test_optimizer.py +++ b/tests/test_utils/test_optimizer.py @@ -2,7 +2,6 @@ import pytest import torch - from phlower.utils import OptimizerSelector @@ -23,7 +22,7 @@ ("MyOptimizer", False), ], ) -def test__exist(name, desired): +def test__exist(name: str, desired: bool): assert OptimizerSelector.exist(name) == desired @@ -43,13 +42,13 @@ def test__exist(name, desired): ("SGD"), ], ) -def test__select(name): +def test__select(name: str): optim = OptimizerSelector.select(name) assert optim.__name__ == name @pytest.mark.parametrize("name", ["MyOptimizer", "BestOptimizer"]) -def test__exist_after_register(name): +def test__exist_after_register(name: str): assert not OptimizerSelector.exist(name) dummy = mock.MagicMock(torch.optim.Optimizer) OptimizerSelector.register(name, dummy) diff --git a/tests/test_utils/test_preprocess.py b/tests/test_utils/test_preprocess.py index bbdb81e..6007659 100644 --- a/tests/test_utils/test_preprocess.py +++ b/tests/test_utils/test_preprocess.py @@ -2,32 +2,31 @@ from hypothesis import assume, given from hypothesis import strategies as st from hypothesis.extra import numpy as ex_np - from phlower.utils import convert_to_dumped, get_registered_scaler_names from phlower.utils.enums import PhlowerScalerName @given(st.sampled_from(PhlowerScalerName)) -def test__get_registered_scaler_names(scaler_name_type): +def test__get_registered_scaler_names(scaler_name_type: PhlowerScalerName): names = get_registered_scaler_names() assert scaler_name_type.value in names @given(st.integers()) -def test__dumped_int_object(val): +def test__dumped_int_object(val: int): actual = convert_to_dumped(val) assert actual == val @given(ex_np.arrays(dtype=np.float64, shape=(3, 5))) -def test__dumped_numpy_dense_array(val): +def test__dumped_numpy_dense_array(val: np.ndarray): assume(~np.any(np.isnan(val))) actual = convert_to_dumped(val) assert actual == val.tolist() @given(ex_np.from_dtype(np.dtype("float"))) -def test__dumped_numpy_scalar(val): +def test__dumped_numpy_scalar(val: np.ndarray): actual = convert_to_dumped(val) if np.isnan(val): assert np.isnan(actual) @@ -36,6 +35,6 @@ def test__dumped_numpy_scalar(val): @given(st.lists(st.floats())) -def test_dumped_sequence_object(val): +def test_dumped_sequence_object(val: list[float]): actual = convert_to_dumped(val) assert actual == val diff --git a/tests/test_utils/test_progress_bar.py b/tests/test_utils/test_progress_bar.py index 25307e2..55ed5e0 100644 --- a/tests/test_utils/test_progress_bar.py +++ b/tests/test_utils/test_progress_bar.py @@ -1,10 +1,9 @@ import pytest - from phlower.utils import PhlowerProgressBar @pytest.mark.parametrize("total", [100, 23, 33]) -def test__is_destroyed_after_iteration_completed(total): +def test__is_destroyed_after_iteration_completed(total: int): pbar = PhlowerProgressBar(total) for _ in range(total): @@ -16,7 +15,7 @@ def test__is_destroyed_after_iteration_completed(total): @pytest.mark.parametrize( "total, desc", [(10, "aaaaa"), (100, "ccccc"), (2, "dddd")] ) -def test__consider_inputs(total, desc): +def test__consider_inputs(total: int, desc: str): pbar = PhlowerProgressBar(total=total, desc=desc) pbar._create_pbar() @@ -25,7 +24,7 @@ def test__consider_inputs(total, desc): @pytest.mark.parametrize("desc", ["aaaa", "loss"]) -def test__consider_desc_when_update(desc): +def test__consider_desc_when_update(desc: str): pbar = PhlowerProgressBar(total=100) pbar.update(1, desc=desc) assert pbar._pbar.desc == desc diff --git a/tests/test_utils/test_schedulers.py b/tests/test_utils/test_schedulers.py index 1adb1f4..3d38748 100644 --- a/tests/test_utils/test_schedulers.py +++ b/tests/test_utils/test_schedulers.py @@ -2,7 +2,6 @@ import pytest import torch - from phlower.utils import SchedulerSelector @@ -27,7 +26,7 @@ ("MyScheduler", False), ], ) -def test__exist(name, desired): +def test__exist(name: str, desired: bool): assert SchedulerSelector.exist(name) == desired @@ -51,13 +50,13 @@ def test__exist(name, desired): ("PolynomialLR"), ], ) -def test__select(name): +def test__select(name: str): scheduler = SchedulerSelector.select(name) assert scheduler.__name__ == name @pytest.mark.parametrize("name", ["MyScheduler", "BestScheduler"]) -def test__exist_after_register(name): +def test__exist_after_register(name: str): assert not SchedulerSelector.exist(name) dummy = mock.MagicMock(torch.optim.Optimizer) SchedulerSelector.register(name, dummy) diff --git a/tests/test_utils/test_timer.py b/tests/test_utils/test_timer.py index 03be030..5b9a5a0 100644 --- a/tests/test_utils/test_timer.py +++ b/tests/test_utils/test_timer.py @@ -1,14 +1,15 @@ from unittest import mock import pytest - from phlower.utils import StopWatch @pytest.mark.parametrize( "offset, times, expected", [(0, [0, 100], 100), (10, [0, 100], 110)] ) -def test__stop_function_elapsed_time(offset, times, expected): +def test__stop_function_elapsed_time( + offset: int, times: list[int], expected: int +) -> None: timer = StopWatch(offset) with mock.patch("time.time", side_effect=times): timer.start() @@ -19,7 +20,9 @@ def test__stop_function_elapsed_time(offset, times, expected): @pytest.mark.parametrize( "offset, times, expected", [(0, [0, 100], 100), (10, [0, 100], 110)] ) -def test__watch_function_elapsed_time(offset, times, expected): +def test__watch_function_elapsed_time( + offset: int, times: list[int], expected: int +) -> None: timer = StopWatch(offset) with mock.patch("time.time", side_effect=times): timer.start() @@ -27,7 +30,7 @@ def test__watch_function_elapsed_time(offset, times, expected): assert timer._start is not None -def test__cannot_start_multiple_times(): +def test__cannot_start_multiple_times() -> None: timer = StopWatch() with pytest.raises(ValueError): timer.start() From 3703f15aece887f3a78b969791d962fa0cf51c42 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 19 Aug 2024 16:11:04 +0900 Subject: [PATCH 51/89] fix lint errors --- tests/test_base/test_dimensions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_base/test_dimensions.py b/tests/test_base/test_dimensions.py index 3f3e547..aeba6eb 100644 --- a/tests/test_base/test_dimensions.py +++ b/tests/test_base/test_dimensions.py @@ -12,7 +12,7 @@ st.floats(allow_nan=False), ) ) -def test__equal_when_same_dimension(dict_data: dict[str.float]): +def test__equal_when_same_dimension(dict_data: dict[str, float]): dimension = PhysicalDimensions(dict_data) other = PhysicalDimensions(dict_data) From d6d246f4f24f60ba7dfedde4cb6f79d7f836bc51 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 19 Aug 2024 16:43:59 +0900 Subject: [PATCH 52/89] fix command to perform lint check --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 1380d3e..a103c1f 100644 --- a/Makefile +++ b/Makefile @@ -29,7 +29,7 @@ e2e_test: .PHONY: lint lint: - poetry run python3 -m ruff check --diff + poetry run python3 -m ruff check --output-format=full poetry run python3 -m ruff format --diff # $(MAKE) mypy From b85e0df478a3f3f15b0c85eb367e7f506ff7f712 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 21 Aug 2024 14:21:47 +0900 Subject: [PATCH 53/89] fix typehint --- src/phlower/_fields/_simulation_field.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/phlower/_fields/_simulation_field.py b/src/phlower/_fields/_simulation_field.py index 7da4c05..0a0a5bb 100644 --- a/src/phlower/_fields/_simulation_field.py +++ b/src/phlower/_fields/_simulation_field.py @@ -1,4 +1,5 @@ import abc +from collections.abc import Iterable from phlower import PhlowerTensor from phlower._base import GraphBatchInfo @@ -25,7 +26,7 @@ def __init__( batch_info = {} self._batch_info = batch_info - def keys(self): + def keys(self) -> Iterable[str]: return self._field_tensors.keys() def __getitem__(self, name: str) -> PhlowerTensor: From be44915fd2605c1b187a2bd1f5793f596b2658c6 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 21 Aug 2024 17:55:24 +0900 Subject: [PATCH 54/89] add restart feature to trainer service --- src/phlower/io/_files/_interface.py | 2 +- src/phlower/io/_model_selector.py | 3 +- src/phlower/services/trainer/_optimizer.py | 13 +- src/phlower/services/trainer/_trainer.py | 85 ++++++-- src/phlower/utils/exceptions.py | 8 + tests/test_services/test_trainer/.gitignore | 1 + .../test_services/test_trainer/data/train.yml | 44 ++++ .../test_trainer/data/train_batch_size.yml | 42 ++++ .../test_trainer/test_trainer.py | 190 ++++++++++++++++++ 9 files changed, 367 insertions(+), 21 deletions(-) create mode 100644 tests/test_services/test_trainer/.gitignore create mode 100644 tests/test_services/test_trainer/data/train.yml create mode 100644 tests/test_services/test_trainer/data/train_batch_size.yml create mode 100644 tests/test_services/test_trainer/test_trainer.py diff --git a/src/phlower/io/_files/_interface.py b/src/phlower/io/_files/_interface.py index e4a872f..434ee87 100644 --- a/src/phlower/io/_files/_interface.py +++ b/src/phlower/io/_files/_interface.py @@ -84,7 +84,7 @@ def to_pathlib_object( ) -> pathlib.Path: if isinstance(path, pathlib.Path): return path - if isinstance(path, IPhlowerYamlFile): + if isinstance(path, IPhlowerBaseFile): return path.file_path if isinstance(path, str): return pathlib.Path(path) diff --git a/src/phlower/io/_model_selector.py b/src/phlower/io/_model_selector.py index 4b1415c..7742be3 100644 --- a/src/phlower/io/_model_selector.py +++ b/src/phlower/io/_model_selector.py @@ -1,6 +1,7 @@ import abc import os import pathlib +from typing import Literal import numpy as np import pandas as pd @@ -15,7 +16,7 @@ def select_snapshot_file( directory: os.PathLike | PhlowerDirectory, - selection_mode: str, + selection_mode: Literal["best", "latest", "train_best", "specified"], **kwards, ) -> IPhlowerCheckpointFile: selector = ModelSelectorBuilder.create(selection_mode) diff --git a/src/phlower/services/trainer/_optimizer.py b/src/phlower/services/trainer/_optimizer.py index 1139726..af16489 100644 --- a/src/phlower/services/trainer/_optimizer.py +++ b/src/phlower/services/trainer/_optimizer.py @@ -53,4 +53,15 @@ def step_scheduler(self): scheduler.step() def state_dict(self) -> dict: - return self._optimizer.state_dict() + return { + "optimizer": self._optimizer.state_dict(), + "schedulers": [ + scheduler.state_dict() for scheduler in self._schedulers + ], + } + + def load_state_dict(self, content: dict) -> None: + self._optimizer.load_state_dict(content["optimizer"]) + + for i, state in enumerate(content["schedulers"]): + self._schedulers[i].load_state_dict(state) diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index 6e61f4e..c50b901 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -8,7 +8,11 @@ from phlower._base import PhlowerTensor from phlower.data import DataLoaderBuilder, LazyPhlowerDataset, LumpedTensorData -from phlower.io import PhlowerCheckpointFile, PhlowerYamlFile +from phlower.io import ( + PhlowerCheckpointFile, + PhlowerYamlFile, + select_snapshot_file, +) from phlower.nn import PhlowerGroupModule from phlower.services.loss_operations import LossCalculator from phlower.services.trainer._optimizer import PhlowerOptimizerWrapper @@ -19,6 +23,8 @@ PhlowerTrainerSetting, ) from phlower.utils import PhlowerProgressBar, StopWatch, get_logger +from phlower.utils.enums import ModelSelectionType +from phlower.utils.exceptions import PhlowerRestartTrainingCompletedError _logger = get_logger(__name__) @@ -62,7 +68,8 @@ def __init__( self._scheduled_optimizer = PhlowerOptimizerWrapper.from_setting( self._trainer_setting, model=self._model ) - self._timer = StopWatch() + self._start_epoch = 0 + self._offset_time = 0.0 def _fix_seed(self, seed: int): random.seed(seed) @@ -77,9 +84,11 @@ def train( disable_dimensions: bool = False, encrypt_key: bytes | None = None, ) -> PhlowerTensor: - self._save_setting(output_directory, encrypt_key=encrypt_key) record_io = LogRecordIO(file_path=output_directory / "log.csv") - record_io.write_header() + if self._start_epoch == 0: + # start_epoch > 0 means that this training is restarted. + self._save_setting(output_directory, encrypt_key=encrypt_key) + record_io.write_header() train_dataset = LazyPhlowerDataset( x_variable_names=self._model_setting.network.get_input_keys(), @@ -107,15 +116,15 @@ def train( ) loss_function = LossCalculator.from_setting(self._trainer_setting) - loss: PhlowerTensor | None = None tqdm.write(record_io.get_header()) - self._timer.start() + _timer = StopWatch(offset=self._offset_time) + _timer.start() _train_batch_pbar = PhlowerProgressBar(total=len(train_dataset)) _val_batch_pbar = PhlowerProgressBar(total=len(validation_dataset)) - for epoch in range(self._trainer_setting.n_epoch): + for epoch in range(self._start_epoch, self._trainer_setting.n_epoch): train_losses: list[float] = [] validation_losses: list[float] = [] @@ -145,6 +154,7 @@ def train( self._model.eval() for val_batch in validation_loader: with torch.no_grad(): + val_batch: LumpedTensorData h = self._model.forward( val_batch.x_data, supports=val_batch.sparse_supports ) @@ -164,11 +174,13 @@ def train( train_loss = np.average(train_losses) validation_loss = np.average(validation_losses) + elapsed_time = _timer.watch() + log_record = LogRecord( epoch=epoch, train_loss=train_loss, validation_loss=validation_loss, - elapsed_time=self._timer.watch(), + elapsed_time=elapsed_time, ) tqdm.write(record_io.to_str(log_record)) @@ -181,18 +193,10 @@ def train( output_directory=output_directory, epoch=epoch, validation_loss=validation_loss, + elapsed_time=elapsed_time, ) return loss.detach() - def _load_state( - self, - target_path: pathlib.Path, - device: str | None = None, - decrypt_key: bytes | None = None, - ) -> None: - # Restore model and optimizer and tqdm - ... - def _save_setting( self, output_directory: pathlib.Path, encrypt_key: bytes | None = None ) -> None: @@ -212,13 +216,15 @@ def _save_checkpoint( output_directory: pathlib.Path, epoch: int, validation_loss: float, + elapsed_time: float, encrypt_key: bytes | None = None, ) -> None: data = { "epoch": epoch, "validation_loss": validation_loss, "model_state_dict": self._model.state_dict(), - "optimizer_state_dict": self._scheduled_optimizer.state_dict(), + "scheduled_optimizer": self._scheduled_optimizer.state_dict(), + "elapsed_time": elapsed_time, } prefix = PhlowerCheckpointFile.get_fixed_prefix() file_basename = f"{prefix}{epoch}" @@ -229,3 +235,46 @@ def _save_checkpoint( encrypt_key=encrypt_key, ) return + + def reinit_for_restart( + self, + restart_directory: pathlib.Path, + device: str | None = None, + decrypt_key: bytes | None = None, + ) -> None: + # Restore model and optimizer and tqdm + snapshot_file = select_snapshot_file( + restart_directory, selection_mode=ModelSelectionType.LATEST.value + ) + checkpoint = self.load_state( + snapshot_file=snapshot_file, device=device, decrypt_key=decrypt_key + ) + + self._start_epoch = int(checkpoint["epoch"]) + 1 + self._offset_time = checkpoint["elapsed_time"] + + if self._trainer_setting.n_epoch == self._start_epoch: + raise PhlowerRestartTrainingCompletedError( + "Checkpoint at last epoch exists. " + "Model to restart has already finished" + ) + + def load_state( + self, + snapshot_file: PhlowerCheckpointFile, + device: str | None = None, + decrypt_key: bytes | None = None, + ) -> dict: + # Restore model and optimizer and tqdm + checkpoint = snapshot_file.load(device=device, decrypt_key=decrypt_key) + + self._model.load_state_dict(checkpoint["model_state_dict"]) + self._scheduled_optimizer.load_state_dict( + checkpoint["scheduled_optimizer"] + ) + + # self.loss = checkpoint['loss'] + _logger.info( + f"{snapshot_file.file_path} is successfully " "loaded for restart." + ) + return checkpoint diff --git a/src/phlower/utils/exceptions.py b/src/phlower/utils/exceptions.py index a7cda1b..21c8a17 100644 --- a/src/phlower/utils/exceptions.py +++ b/src/phlower/utils/exceptions.py @@ -43,6 +43,14 @@ class PhlowerFeatureStoreOverwriteError(ValueError): ... +class PhlowerRestartTrainingCompletedError(ValueError): + """This error raises when trying to restart train job + which has already completed. + """ + + ... + + class PhlowerSparseUnsupportedError(ValueError): """ This error raises when trying to call methods not supported for sparse diff --git a/tests/test_services/test_trainer/.gitignore b/tests/test_services/test_trainer/.gitignore new file mode 100644 index 0000000..6253c90 --- /dev/null +++ b/tests/test_services/test_trainer/.gitignore @@ -0,0 +1 @@ +_out \ No newline at end of file diff --git a/tests/test_services/test_trainer/data/train.yml b/tests/test_services/test_trainer/data/train.yml new file mode 100644 index 0000000..0f8fa56 --- /dev/null +++ b/tests/test_services/test_trainer/data/train.yml @@ -0,0 +1,44 @@ + +training: + batch_size: 1 + random_seed: 0 + n_epoch: 10 + loss_setting: + name2loss: + nodal_last_u: "mse" + optimizer_setting: + optimizer: SGD + parameters: + lr: 0.0001 + + +model: + variable_dimensions: + nodal_initial_u: {"L": 1, "T": -1} + nodal_last_u: {"L": 1, "T": -1} + nodal_nadj: {} + + network: + nn_type: GROUP + name: DEMO + inputs: + - name: nodal_initial_u + n_dim: 1 + + outputs: + - name: nodal_last_u + n_dim: 1 + + support_names: ["nodal_nadj"] + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - nodal_initial_u + output_key: nodal_last_u + nn_parameters: + nodes: [-1, 16, 1] + activations: ["identity", "identity"] + bias: False + diff --git a/tests/test_services/test_trainer/data/train_batch_size.yml b/tests/test_services/test_trainer/data/train_batch_size.yml new file mode 100644 index 0000000..409c59e --- /dev/null +++ b/tests/test_services/test_trainer/data/train_batch_size.yml @@ -0,0 +1,42 @@ + +training: + batch_size: 3 + random_seed: 0 + optimizer_setting: + optimizer: SGD + parameters: + lr: 0.0001 + loss_setting: + name2loss: + nodal_last_u: "mse" + +model: + variable_dimensions: + nodal_initial_u: {"L": 1, "T": -1} + nodal_last_u: {"L": 1, "T": -1} + nodal_nadj: {} + + network: + nn_type: GROUP + name: DEMO + inputs: + - name: nodal_initial_u + n_dim: 1 + + outputs: + - name: nodal_last_u + n_dim: 1 + + support_names: ["nodal_nadj"] + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - nodal_initial_u + output_key: nodal_last_u + nn_parameters: + nodes: [-1, 16, 1] + activations: ["identity", "identity"] + bias: False + diff --git a/tests/test_services/test_trainer/test_trainer.py b/tests/test_services/test_trainer/test_trainer.py new file mode 100644 index 0000000..bbad967 --- /dev/null +++ b/tests/test_services/test_trainer/test_trainer.py @@ -0,0 +1,190 @@ +import pathlib +import random +import shutil +from collections.abc import Callable + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp +import torch +import yaml +from phlower import PhlowerTensor +from phlower.io import PhlowerDirectory, select_snapshot_file +from phlower.services.trainer import PhlowerTrainer +from phlower.settings import PhlowerSetting +from phlower.utils.exceptions import PhlowerRestartTrainingCompletedError + +_OUTPUT_DIR = pathlib.Path(__file__).parent / "_out" +_SETTINGS_DIR = pathlib.Path(__file__).parent / "data" + + +@pytest.fixture(scope="module") +def prepare_sample_preprocessed_files(): + random.seed(11) + np.random.seed(11) + if _OUTPUT_DIR.exists(): + shutil.rmtree(_OUTPUT_DIR) + _OUTPUT_DIR.mkdir() + + base_preprocessed_dir = _OUTPUT_DIR / "preprocessed" + base_preprocessed_dir.mkdir() + + n_cases = 3 + dtype = np.float32 + for i in range(n_cases): + n_nodes = 100 * (i + 1) + preprocessed_dir = base_preprocessed_dir / f"case_{i}" + preprocessed_dir.mkdir() + + nodal_initial_u = np.random.rand(n_nodes, 3, 1) + np.save( + preprocessed_dir / "nodal_initial_u.npy", + nodal_initial_u.astype(dtype), + ) + + # nodal_last_u = np.random.rand(n_nodes, 3, 1) + np.save( + preprocessed_dir / "nodal_last_u.npy", nodal_initial_u.astype(dtype) + ) + + rng = np.random.default_rng() + nodal_nadj = sp.random(n_nodes, n_nodes, density=0.1, random_state=rng) + sp.save_npz( + preprocessed_dir / "nodal_nadj", nodal_nadj.tocoo().astype(dtype) + ) + + (preprocessed_dir / "preprocessed").touch() + + +@pytest.fixture(scope="module") +def simple_training(prepare_sample_preprocessed_files: None) -> PhlowerTensor: + phlower_path = PhlowerDirectory(_OUTPUT_DIR) + + preprocessed_directories = list( + phlower_path.find_directory( + required_filename="preprocessed", recursive=True + ) + ) + + setting = PhlowerSetting.read_yaml(_SETTINGS_DIR / "train.yml") + + trainer = PhlowerTrainer.from_setting(setting) + output_directory = _OUTPUT_DIR / "model" + if output_directory.exists(): + shutil.rmtree(output_directory) + + loss = trainer.train( + train_directories=preprocessed_directories, + validation_directories=preprocessed_directories, + output_directory=output_directory, + ) + return loss + + +@pytest.mark.e2e_test +def test__training_with_multiple_batch_size( + prepare_sample_preprocessed_files: None, +): + phlower_path = PhlowerDirectory(_OUTPUT_DIR) + preprocessed_directories = list( + phlower_path.find_directory( + required_filename="preprocessed", recursive=True + ) + ) + + setting = PhlowerSetting.read_yaml(_SETTINGS_DIR / "train_batch_size.yml") + assert setting.training.batch_size > 1 + + trainer = PhlowerTrainer.from_setting(setting) + output_directory = _OUTPUT_DIR / "model_batch_size" + if output_directory.exists(): + shutil.rmtree(output_directory) + + loss = trainer.train( + train_directories=preprocessed_directories, + validation_directories=preprocessed_directories, + output_directory=output_directory, + ) + assert loss.has_dimension + assert not torch.isinf(loss.to_tensor()) + assert not torch.isnan(loss.to_tensor()) + + +@pytest.mark.e2e_test +def test__simple_training(simple_training: PhlowerTensor): + loss: PhlowerTensor = simple_training + + assert loss.has_dimension + assert not torch.isinf(loss.to_tensor()) + assert not torch.isnan(loss.to_tensor()) + + +@pytest.fixture +def perform_restart() -> Callable[[int | None], None]: + def restart_training(n_epoch: int | None): + if n_epoch is not None: + # NOTE: overwrite n_epoch to restart + with open(_SETTINGS_DIR / "train.yml") as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + content["training"]["n_epoch"] = n_epoch + + setting = PhlowerSetting(**content) + trainer = PhlowerTrainer.from_setting(setting) + + restart_directory = _OUTPUT_DIR / "model" + trainer.reinit_for_restart(restart_directory=restart_directory) + + phlower_path = PhlowerDirectory(_OUTPUT_DIR) + preprocessed_directories = list( + phlower_path.find_directory( + required_filename="preprocessed", recursive=True + ) + ) + trainer.train( + output_directory=restart_directory, + train_directories=preprocessed_directories, + validation_directories=preprocessed_directories, + ) + + return restart_training + + +def test__not_allowed_restart_when_all_epoch_is_finished( + simple_training: PhlowerTensor, + perform_restart: Callable[[int | None], None], +): + with pytest.raises(PhlowerRestartTrainingCompletedError): + perform_restart() + + +def test__last_epoch_is_update_after_restart( + simple_training: PhlowerTensor, + perform_restart: Callable[[int | None], None], +): + last_snapshot = select_snapshot_file(_OUTPUT_DIR / "model", "latest") + assert last_snapshot.file_path.name.startswith("snapshot_epoch_9") + + n_epoch = 12 + perform_restart(n_epoch) + + last_snapshot = select_snapshot_file(_OUTPUT_DIR / "model", "latest") + assert last_snapshot.file_path.name.startswith( + f"snapshot_epoch_{n_epoch - 1}" + ) + + # check log.csv + df = pd.read_csv( + _OUTPUT_DIR / "model/log.csv", + header=0, + index_col=None, + skipinitialspace=True, + ) + + assert max(df.loc[:, "epoch"]) == n_epoch - 1 + + # check elapsed_time increases monotonically + _prev = 0.0 + for v in df.loc[:, "elapsed_time"]: + assert v > _prev + _prev = v From b0d29cda5f3373e8de37e1466dbb81675d11676a Mon Sep 17 00:00:00 2001 From: sakamoto Date: Thu, 22 Aug 2024 11:48:41 +0900 Subject: [PATCH 55/89] fix unit tests for trainer --- tests/test_services/test_trainer/test_optimizer.py | 13 +++++++------ tests/test_services/test_trainer/test_trainer.py | 11 +++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/test_services/test_trainer/test_optimizer.py b/tests/test_services/test_trainer/test_optimizer.py index 362b8cc..51f37ef 100644 --- a/tests/test_services/test_trainer/test_optimizer.py +++ b/tests/test_services/test_trainer/test_optimizer.py @@ -52,7 +52,7 @@ def test__pass_kwargs_when_call_from_setting( @pytest.mark.parametrize( - "optimizer, lr, weight_decay, desired_optimizer", + "optimizer_name, lr, weight_decay, desired_optimizer", [ ("Adam", 0.001, 0, torch.optim.Adam), ("Adam", 0.0003, 0, torch.optim.Adam), @@ -60,7 +60,7 @@ def test__pass_kwargs_when_call_from_setting( ], ) def test__optimizer_parameters( - optimizer: str, + optimizer_name: str, lr: float, weight_decay: float, desired_optimizer: type[torch.optim.Optimizer], @@ -68,7 +68,7 @@ def test__optimizer_parameters( model = torch.nn.Linear(in_features=10, out_features=10) optimizer = PhlowerOptimizerWrapper( parameters=model.parameters(), - optimizer=optimizer, + optimizer=optimizer_name, optimizer_kwargs={"lr": lr, "weight_decay": weight_decay}, schedulers={}, ) @@ -76,9 +76,10 @@ def test__optimizer_parameters( assert isinstance(optimizer._optimizer, desired_optimizer) state_dict = optimizer.state_dict() - assert len(state_dict["param_groups"]) == 1 - assert state_dict["param_groups"][0]["lr"] == lr - assert state_dict["param_groups"][0]["weight_decay"] == weight_decay + params = state_dict["optimizer"]["param_groups"] + assert len(params) == 1 + assert params[0]["lr"] == lr + assert params[0]["weight_decay"] == weight_decay @pytest.mark.parametrize( diff --git a/tests/test_services/test_trainer/test_trainer.py b/tests/test_services/test_trainer/test_trainer.py index bbad967..c56bd52 100644 --- a/tests/test_services/test_trainer/test_trainer.py +++ b/tests/test_services/test_trainer/test_trainer.py @@ -82,7 +82,6 @@ def simple_training(prepare_sample_preprocessed_files: None) -> PhlowerTensor: return loss -@pytest.mark.e2e_test def test__training_with_multiple_batch_size( prepare_sample_preprocessed_files: None, ): @@ -111,7 +110,6 @@ def test__training_with_multiple_batch_size( assert not torch.isnan(loss.to_tensor()) -@pytest.mark.e2e_test def test__simple_training(simple_training: PhlowerTensor): loss: PhlowerTensor = simple_training @@ -122,12 +120,13 @@ def test__simple_training(simple_training: PhlowerTensor): @pytest.fixture def perform_restart() -> Callable[[int | None], None]: - def restart_training(n_epoch: int | None): + def restart_training(n_epoch: int | None = None): + with open(_SETTINGS_DIR / "train.yml") as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + if n_epoch is not None: # NOTE: overwrite n_epoch to restart - with open(_SETTINGS_DIR / "train.yml") as fr: - content = yaml.load(fr, Loader=yaml.SafeLoader) - content["training"]["n_epoch"] = n_epoch + content["training"]["n_epoch"] = n_epoch setting = PhlowerSetting(**content) trainer = PhlowerTrainer.from_setting(setting) From 4bdef42a878217ba901601b3e6f95030e05882a7 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Thu, 22 Aug 2024 12:19:47 +0900 Subject: [PATCH 56/89] fix setting file --- tests/test_services/test_trainer/data/train.yml | 6 ++++-- tests/test_services/test_trainer/data/train_batch_size.yml | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_services/test_trainer/data/train.yml b/tests/test_services/test_trainer/data/train.yml index 0f8fa56..07e2245 100644 --- a/tests/test_services/test_trainer/data/train.yml +++ b/tests/test_services/test_trainer/data/train.yml @@ -18,6 +18,10 @@ model: nodal_last_u: {"L": 1, "T": -1} nodal_nadj: {} + fields: + field_names: + - "nodal_nadj" + network: nn_type: GROUP name: DEMO @@ -29,8 +33,6 @@ model: - name: nodal_last_u n_dim: 1 - support_names: ["nodal_nadj"] - modules: - nn_type: MLP name: MLP0 diff --git a/tests/test_services/test_trainer/data/train_batch_size.yml b/tests/test_services/test_trainer/data/train_batch_size.yml index 409c59e..c720955 100644 --- a/tests/test_services/test_trainer/data/train_batch_size.yml +++ b/tests/test_services/test_trainer/data/train_batch_size.yml @@ -16,6 +16,10 @@ model: nodal_last_u: {"L": 1, "T": -1} nodal_nadj: {} + fields: + field_names: + - "nodal_nadj" + network: nn_type: GROUP name: DEMO @@ -27,8 +31,6 @@ model: - name: nodal_last_u n_dim: 1 - support_names: ["nodal_nadj"] - modules: - nn_type: MLP name: MLP0 From 168e1967af25674f1b5e20effbe74733bc60d029 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Wed, 25 Sep 2024 14:02:40 +0900 Subject: [PATCH 57/89] add tensor pattern object --- src/phlower/_base/tensors/__init__.py | 1 + src/phlower/_base/tensors/_interface.py | 19 +++- src/phlower/_base/tensors/_phlower_tensor.py | 87 +++++-------------- src/phlower/_base/tensors/_tensor_shape.py | 71 +++++++++++++++ src/phlower/nn/_core_modules/_functions.py | 9 +- .../test_tensors/test__phlower_tensor.py | 7 +- .../test_tensors/test_tensor_shape.py | 69 +++++++++++++++ 7 files changed, 190 insertions(+), 73 deletions(-) create mode 100644 src/phlower/_base/tensors/_tensor_shape.py create mode 100644 tests/test_base/test_tensors/test_tensor_shape.py diff --git a/src/phlower/_base/tensors/__init__.py b/src/phlower/_base/tensors/__init__.py index 02f95d8..44d3615 100644 --- a/src/phlower/_base/tensors/__init__.py +++ b/src/phlower/_base/tensors/__init__.py @@ -3,3 +3,4 @@ phlower_dimension_tensor, ) from phlower._base.tensors._phlower_tensor import PhlowerTensor, phlower_tensor +from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern diff --git a/src/phlower/_base/tensors/_interface.py b/src/phlower/_base/tensors/_interface.py index a9a429c..ce3d8ee 100644 --- a/src/phlower/_base/tensors/_interface.py +++ b/src/phlower/_base/tensors/_interface.py @@ -6,6 +6,7 @@ import torch from phlower._base.tensors import PhlowerDimensionTensor +from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern class IPhlowerTensor(metaclass=abc.ABCMeta): @@ -21,6 +22,10 @@ def dimension(self) -> PhlowerDimensionTensor | None: ... @abc.abstractmethod def shape(self) -> torch.Size: ... + @property + @abc.abstractmethod + def shape_pattern(self) -> PhlowerTensorShapePattern: ... + @property @abc.abstractmethod def is_sparse(self) -> bool: ... @@ -72,7 +77,19 @@ def to(self, device: str, non_blocking: bool = False) -> None: ... @abc.abstractmethod def backward(self) -> None: ... - @abc.abstractclassmethod + @abc.abstractmethod + def to_vertexwise(self) -> tuple[IPhlowerTensor, str]: ... + + @abc.abstractmethod + def rearrange( + self, + pattern: str, + is_time_series: bool = False, + is_voxel: bool = False, + **kwargs: dict[str, int], + ) -> IPhlowerTensor: ... + + @abc.abstractmethod def __torch_function__( cls, func: Callable, diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 0b5655a..06014bc 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -14,6 +14,7 @@ phlower_dimension_tensor, ) from phlower._base.tensors._interface import IPhlowerTensor +from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern from phlower._base.tensors._unsupported_function_names import ( UNSUPPORTED_FUNCTION_NAMES, ) @@ -110,8 +111,9 @@ def __init__( ) self._tensor = tensor self._dimension_tensor = dimension_tensor - self._is_time_series = is_time_series - self._is_voxel = is_voxel + self._shape_pattern = PhlowerTensorShapePattern( + self._tensor.shape, is_time_series=is_time_series, is_voxel=is_voxel + ) @property def has_dimension(self) -> bool: @@ -123,7 +125,11 @@ def dimension(self) -> PhlowerDimensionTensor | None: @property def shape(self) -> torch.Size: - return self._tensor.shape + return self._shape_pattern.shape + + @property + def shape_pattern(self) -> PhlowerTensorShapePattern: + return self._shape_pattern @property def is_sparse(self) -> bool: @@ -131,11 +137,11 @@ def is_sparse(self) -> bool: @property def is_time_series(self) -> bool: - return self._is_time_series + return self._shape_pattern.is_time_series @property def is_voxel(self) -> bool: - return self._is_voxel + return self._shape_pattern.is_voxel def __str__(self) -> str: return ( @@ -219,13 +225,8 @@ def rank(self) -> int: raise PhlowerSparseUnsupportedError( "Cannot call rank() for sparse PhlowerTensor" ) - size = self.size() - start = 1 - if self.is_time_series: - start += 1 - if self.is_voxel: - start += 2 - return len(size[start:-1]) + + return self._shape_pattern.rank_size def n_vertices(self) -> int: """Returns the number of vertices.""" @@ -233,15 +234,8 @@ def n_vertices(self) -> int: raise PhlowerSparseUnsupportedError( "Cannot call n_vertices() for sparse PhlowerTensor" ) - size = self.size() - start = 0 - if self.is_time_series: - start += 1 - - if self.is_voxel: - return np.prod(size[start : start + 3]) - return size[start] + return self._shape_pattern.n_vertices() def indices(self) -> torch.Tensor: return self._tensor.indices() @@ -249,55 +243,22 @@ def indices(self) -> torch.Tensor: def values(self) -> torch.Tensor: return self._tensor.values() - def to_vertexwise(self) -> PhlowerTensor: + def to_vertexwise(self) -> tuple[PhlowerTensor, str]: """ Convert to vertexwise 2D tensor which has (n_vertices, -1) shape. Returns: vertexwise_tensor : PhlowerTensor Vertexwise PhlowerTensor object. - original_pattern : str - Pattern of the original shape. Can be used for rearrange. resultant_pattern : str Pattern of the resultant shape. Can be used for rearrange. - dict_shape : dict[str, int] - Dict of original shape. Can be used for rearrange. """ - shape = self.shape - dict_shape = {} - - space_start = 0 - if self.is_time_series: - t_pattern = "t " - space_start += 1 - dict_shape.update({"t": shape[0]}) - else: - t_pattern = "" - if self.is_voxel: - space_pattern = "x y z " - dict_shape.update( - { - "x": shape[space_start], - "y": shape[space_start + 1], - "z": shape[space_start + 2], - } - ) - feat_start = space_start + 3 - else: - space_pattern = "n " - dict_shape.update({"n": shape[space_start]}) - feat_start = space_start + 1 - - feat_pattern = " ".join( - [f"a{i}" for i in range(len(shape[feat_start:]))] + original_pattern = self._shape_pattern.pattern + resultant_pattern = ( + f"({self._shape_pattern.space_pattern}) " + f"({self._shape_pattern.time_series_pattern} " + f"{self._shape_pattern.feature_pattern})" ) - # Do not include the last axis in case modified by NNs. - dict_shape.update( - {f"a{i}": s for i, s in enumerate(shape[feat_start:-1])} - ) - - original_pattern = f"{t_pattern}{space_pattern}{feat_pattern}" - resultant_pattern = f"({space_pattern}) ({t_pattern}{feat_pattern})" tensor_2d = einops.rearrange( self.to_tensor(), f"{original_pattern} -> {resultant_pattern}" ) @@ -308,9 +269,7 @@ def to_vertexwise(self) -> PhlowerTensor: is_time_series=False, is_voxel=False, ), - original_pattern, resultant_pattern, - dict_shape, ) def rearrange( @@ -318,7 +277,7 @@ def rearrange( pattern: str, is_time_series: bool = False, is_voxel: bool = False, - **kwargs, + **kwargs: dict[str, int], ) -> PhlowerTensor: tensor = self.to_tensor() rearranged = einops.rearrange(tensor, pattern, **kwargs) @@ -379,9 +338,9 @@ def __torch_function__( # NOTE: Assume flags for the first tensor is preserved is_time_series = _recursive_resolve( - args, "_is_time_series", return_first_only=True + args, "is_time_series", return_first_only=True ) - is_voxel = _recursive_resolve(args, "_is_voxel", return_first_only=True) + is_voxel = _recursive_resolve(args, "is_voxel", return_first_only=True) if not _has_dimension(args): # Unit calculation is not considered when unit tensor is not found. diff --git a/src/phlower/_base/tensors/_tensor_shape.py b/src/phlower/_base/tensors/_tensor_shape.py new file mode 100644 index 0000000..1d73c6e --- /dev/null +++ b/src/phlower/_base/tensors/_tensor_shape.py @@ -0,0 +1,71 @@ +import functools + +import numpy as np +import torch + + +class PhlowerTensorShapePattern: + def __init__(self, shape: torch.Size, is_time_series: bool, is_voxel: bool): + self._shape = shape + self._is_time_series = is_time_series + self._is_voxel = is_voxel + + def pattern_to_ndim(self, drop_last: bool = False) -> dict[str, int]: + chars = self.pattern.split(" ") + if drop_last: + chars.pop() + + return {c: self._shape[i] for i, c in enumerate(chars)} + + def n_vertices(self) -> int: + start = 1 if self._is_time_series else 0 + + if self._is_voxel: + return np.prod(self._shape[start : start + 3]) + return self._shape[start] + + @property + def is_time_series(self) -> bool: + return self._is_time_series + + @property + def is_voxel(self) -> bool: + return self._is_voxel + + @property + def shape(self) -> torch.Size: + return self._shape + + @property + def rank_size(self) -> int: + return len(self._shape[self.feature_start_dim : -1]) + + @functools.cached_property + def pattern(self) -> str: + patterns = [ + self.time_series_pattern, + self.space_pattern, + self.feature_pattern, + ] + + return " ".join([p for p in patterns if len(p) != 0]) + + @property + def time_series_pattern(self) -> str: + return "t" if self._is_time_series else "" + + @property + def space_pattern(self) -> str: + return "x y z" if self._is_voxel else "n" + + @property + def feature_start_dim(self) -> int: + offset_time = 1 if self._is_time_series else 0 + offset_space = 3 if self._is_voxel else 1 + + return offset_time + offset_space + + @functools.cached_property + def feature_pattern(self) -> str: + start = self.feature_start_dim + return " ".join([f"a{i}" for i in range(len(self._shape[start:]))]) diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 5f14f6f..9173036 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -68,15 +68,16 @@ def spmm( IPhlowerTensor: Resultant tensor. """ - h, original_pattern, resultant_pattern, dict_shape = x.to_vertexwise() - pattern = f"{resultant_pattern} -> {original_pattern}" + h, resultant_pattern = x.to_vertexwise() + restore_pattern = f"{resultant_pattern} -> {x.shape_pattern.pattern}" + restore_ndim = x.shape_pattern.pattern_to_ndim(drop_last=True) for _ in range(repeat): h = torch.sparse.mm(sparse, h) return h.rearrange( - pattern, + restore_pattern, is_time_series=x.is_time_series, is_voxel=x.is_voxel, - **dict_shape, + **restore_ndim, ) diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test__phlower_tensor.py index 360cf70..0e28beb 100644 --- a/tests/test_base/test_tensors/test__phlower_tensor.py +++ b/tests/test_base/test_tensors/test__phlower_tensor.py @@ -265,11 +265,10 @@ def test__to_vertexwise_inverse( phlower_tensor = PhlowerTensor( torch_tensor, is_time_series=is_time_series, is_voxel=is_voxel ) - vertexwise, original_pattern, resultant_pattern, dict_shape = ( - phlower_tensor.to_vertexwise() - ) + vertexwise, resultant_pattern = phlower_tensor.to_vertexwise() assert len(vertexwise.shape) == 2 - pattern = f"{resultant_pattern} -> {original_pattern}" + pattern = f"{resultant_pattern} -> {phlower_tensor.shape_pattern.pattern}" + dict_shape = phlower_tensor.shape_pattern.pattern_to_ndim(drop_last=True) actual = vertexwise.rearrange( pattern, is_time_series=is_time_series, is_voxel=is_voxel, **dict_shape ) diff --git a/tests/test_base/test_tensors/test_tensor_shape.py b/tests/test_base/test_tensors/test_tensor_shape.py new file mode 100644 index 0000000..05e0809 --- /dev/null +++ b/tests/test_base/test_tensors/test_tensor_shape.py @@ -0,0 +1,69 @@ +import pytest +import torch +from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern + + +@pytest.mark.parametrize( + "shapes, desired", + [((5, 6, 7, 9, 10), "n a0 a1 a2 a3"), ((5, 10), "n a0"), ((5,), "n")], +) +def test__detect_not_time_series_vertexwise_pattern( + shapes: tuple[int], desired: str +): + shape = PhlowerTensorShapePattern( + torch.Size(shapes), is_time_series=False, is_voxel=False + ) + + assert shape.pattern == desired + + +@pytest.mark.parametrize( + "shapes, desired", + [ + ((5, 6, 7, 9, 10), "t n a0 a1 a2"), + ((5, 10), "t n"), + ((100, 5, 23), "t n a0"), + ], +) +def test__detect_time_series_vertexwise_pattern( + shapes: tuple[int], desired: str +): + shape = PhlowerTensorShapePattern( + torch.Size(shapes), is_time_series=True, is_voxel=False + ) + + assert shape.pattern == desired + + +@pytest.mark.parametrize( + "shapes, desired", + [ + ((5, 6, 7, 9, 10), "x y z a0 a1"), + ((12, 6, 7, 9), "x y z a0"), + ((5, 6, 7, 9, 10, 78, 100), "x y z a0 a1 a2 a3"), + ], +) +def test__detect_not_time_series_voxel_pattern( + shapes: tuple[int], desired: str +): + shape = PhlowerTensorShapePattern( + torch.Size(shapes), is_time_series=False, is_voxel=True + ) + + assert shape.pattern == desired + + +@pytest.mark.parametrize( + "shapes, desired", + [ + ((5, 6, 7, 9, 10), "t x y z a0"), + ((12, 6, 7, 9), "t x y z"), + ((5, 6, 7, 9, 10, 78, 100), "t x y z a0 a1 a2"), + ], +) +def test__detect_time_series_voxel_pattern(shapes: tuple[int], desired: str): + shape = PhlowerTensorShapePattern( + torch.Size(shapes), is_time_series=True, is_voxel=True + ) + + assert shape.pattern == desired From 3c7864bfcfd3e45dcae65b26e9841481591e1187 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Thu, 26 Sep 2024 16:53:31 +0900 Subject: [PATCH 58/89] fix test file name --- .../test_tensors/{test__dimensions.py => test_dimensions.py} | 0 .../{test__phlower_tensor.py => test_phlower_tensor.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/test_base/test_tensors/{test__dimensions.py => test_dimensions.py} (100%) rename tests/test_base/test_tensors/{test__phlower_tensor.py => test_phlower_tensor.py} (100%) diff --git a/tests/test_base/test_tensors/test__dimensions.py b/tests/test_base/test_tensors/test_dimensions.py similarity index 100% rename from tests/test_base/test_tensors/test__dimensions.py rename to tests/test_base/test_tensors/test_dimensions.py diff --git a/tests/test_base/test_tensors/test__phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py similarity index 100% rename from tests/test_base/test_tensors/test__phlower_tensor.py rename to tests/test_base/test_tensors/test_phlower_tensor.py From caf4767b19ab0154f2e5238bd206913a037b6f0c Mon Sep 17 00:00:00 2001 From: sakamoto Date: Fri, 27 Sep 2024 18:13:07 +0900 Subject: [PATCH 59/89] upgrade shape pattern object to handle pattern string --- Makefile | 2 +- src/phlower/_base/tensors/__init__.py | 2 +- .../_base/tensors/_dimension_tensor.py | 11 +- src/phlower/_base/tensors/_interface.py | 9 +- src/phlower/_base/tensors/_phlower_tensor.py | 129 ++++++++++----- src/phlower/_base/tensors/_shape_operation.py | 1 + src/phlower/_base/tensors/_tensor_shape.py | 154 +++++++++++++++--- src/phlower/nn/_core_modules/_functions.py | 106 ++++-------- .../test_tensors/test_phlower_tensor.py | 14 +- .../test_tensors/test_tensor_shape.py | 18 +- tests/test_nn/test_core_modules/test_gcn.py | 2 +- 11 files changed, 292 insertions(+), 156 deletions(-) create mode 100644 src/phlower/_base/tensors/_shape_operation.py diff --git a/Makefile b/Makefile index a103c1f..c527355 100644 --- a/Makefile +++ b/Makefile @@ -14,8 +14,8 @@ mypy: .PHONY: format format: - poetry run python3 -m ruff check --fix poetry run python3 -m ruff format + poetry run python3 -m ruff check --fix .PHONY: test test: diff --git a/src/phlower/_base/tensors/__init__.py b/src/phlower/_base/tensors/__init__.py index 44d3615..7ed3b37 100644 --- a/src/phlower/_base/tensors/__init__.py +++ b/src/phlower/_base/tensors/__init__.py @@ -3,4 +3,4 @@ phlower_dimension_tensor, ) from phlower._base.tensors._phlower_tensor import PhlowerTensor, phlower_tensor -from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern +from phlower._base.tensors._tensor_shape import PhlowerShapePattern diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index af3bbc5..67be76c 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -329,7 +329,7 @@ def concatenate( @dimension_wrap_implements(torch.tanh) def tanh(tensor: PhlowerDimensionTensor) -> PhlowerDimensionTensor: - if not tensor.is_dimensionless: + if not tensor.is_dimensionless(): raise DimensionIncompatibleError( f"Should be dimensionless to apply tanh but {tensor}" ) @@ -340,8 +340,15 @@ def tanh(tensor: PhlowerDimensionTensor) -> PhlowerDimensionTensor: def leaky_relu( tensor: PhlowerDimensionTensor, *args: Any, **kwargs: Any ) -> PhlowerDimensionTensor: - if not tensor.is_dimensionless: + if not tensor.is_dimensionless(): raise DimensionIncompatibleError( f"Should be dimensionless to apply leaky_relu but {tensor}" ) return tensor + + +@dimension_wrap_implements(torch.unsqueeze) +def unsqueeze( + input: PhlowerDimensionTensor, dim: int +) -> PhlowerDimensionTensor: + return input diff --git a/src/phlower/_base/tensors/_interface.py b/src/phlower/_base/tensors/_interface.py index ce3d8ee..513182e 100644 --- a/src/phlower/_base/tensors/_interface.py +++ b/src/phlower/_base/tensors/_interface.py @@ -6,7 +6,7 @@ import torch from phlower._base.tensors import PhlowerDimensionTensor -from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern +from phlower._base.tensors._tensor_shape import PhlowerShapePattern class IPhlowerTensor(metaclass=abc.ABCMeta): @@ -24,7 +24,7 @@ def shape(self) -> torch.Size: ... @property @abc.abstractmethod - def shape_pattern(self) -> PhlowerTensorShapePattern: ... + def shape_pattern(self) -> PhlowerShapePattern: ... @property @abc.abstractmethod @@ -84,11 +84,12 @@ def to_vertexwise(self) -> tuple[IPhlowerTensor, str]: ... def rearrange( self, pattern: str, - is_time_series: bool = False, - is_voxel: bool = False, **kwargs: dict[str, int], ) -> IPhlowerTensor: ... + @abc.abstractmethod + def as_pattern(self, pattern: str) -> IPhlowerTensor: ... + @abc.abstractmethod def __torch_function__( cls, diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 06014bc..ac79d45 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Sequence -from typing import Any +from typing import Any, TypeAlias, overload import einops import numpy as np @@ -14,7 +14,7 @@ phlower_dimension_tensor, ) from phlower._base.tensors._interface import IPhlowerTensor -from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern +from phlower._base.tensors._tensor_shape import PhlowerShapePattern from phlower._base.tensors._unsupported_function_names import ( UNSUPPORTED_FUNCTION_NAMES, ) @@ -26,33 +26,60 @@ PhlowerUnsupportedTorchFunctionError, ) +PhysicDimensionLikeObject: TypeAlias = ( + PhysicalDimensions + | PhlowerDimensionTensor + | torch.Tensor + | dict[str, float] + | list[float] + | tuple[float] +) + + logger = get_logger(__name__) +@overload def phlower_tensor( tensor: torch.Tensor | PhlowerTensor, - dimension: ( - PhysicalDimensions - | PhlowerDimensionTensor - | torch.Tensor - | dict[str, float] - | list[float] - | tuple[float] - | None - ) = None, + dimension: PhysicDimensionLikeObject | None = None, is_time_series: bool = False, is_voxel: bool = False, -) -> PhlowerTensor: +) -> PhlowerTensor: ... + + +@overload +def phlower_tensor( + tensor: torch.Tensor | PhlowerTensor, + dimension: PhysicDimensionLikeObject | None = None, + pattern: str = "n...", +) -> PhlowerTensor: ... + + +def phlower_tensor( + tensor: torch.Tensor | PhlowerTensor, + dimension: PhysicDimensionLikeObject | None = None, + is_time_series: bool | None = None, + is_voxel: bool | None = None, + pattern: str | None = None, +): if isinstance(tensor, PhlowerTensor): if dimension is not None: - logger.warning( - "Input dimension_tensor and sparse_batch_info are ignored." - ) - + logger.warning("Input dimension_tensor are ignored.") return tensor dimension_tensor = _resolve_dimension_arg(dimension) + if pattern is not None: + if (is_time_series is not None) or (is_voxel is not None): + raise ValueError( + "pattern is not allowed to be used " + "with is_time_series and is_voxel " + ) + return PhlowerTensor.from_pattern( + tensor, dimension_tensor=dimension_tensor, pattern=pattern + ) + return PhlowerTensor( tensor=tensor, dimension_tensor=dimension_tensor, @@ -98,6 +125,24 @@ class PhlowerTensor(IPhlowerTensor): """ + @classmethod + def from_pattern( + cls, + tensor: torch.Tensor, + dimension_tensor: PhlowerDimensionTensor | None = None, + pattern: str | None = None, + ) -> PhlowerTensor: + phlower_shape: PhlowerShapePattern = PhlowerShapePattern.from_pattern( + tensor.shape, pattern + ) + + return PhlowerTensor( + tensor=tensor, + dimension_tensor=dimension_tensor, + is_time_series=phlower_shape.is_time_series, + is_voxel=phlower_shape.is_voxel, + ) + def __init__( self, tensor: torch.Tensor, @@ -111,7 +156,7 @@ def __init__( ) self._tensor = tensor self._dimension_tensor = dimension_tensor - self._shape_pattern = PhlowerTensorShapePattern( + self._phlower_shape = PhlowerShapePattern( self._tensor.shape, is_time_series=is_time_series, is_voxel=is_voxel ) @@ -125,11 +170,11 @@ def dimension(self) -> PhlowerDimensionTensor | None: @property def shape(self) -> torch.Size: - return self._shape_pattern.shape + return self._phlower_shape.shape @property - def shape_pattern(self) -> PhlowerTensorShapePattern: - return self._shape_pattern + def shape_pattern(self) -> PhlowerShapePattern: + return self._phlower_shape @property def is_sparse(self) -> bool: @@ -137,11 +182,11 @@ def is_sparse(self) -> bool: @property def is_time_series(self) -> bool: - return self._shape_pattern.is_time_series + return self._phlower_shape.is_time_series @property def is_voxel(self) -> bool: - return self._shape_pattern.is_voxel + return self._phlower_shape.is_voxel def __str__(self) -> str: return ( @@ -226,7 +271,7 @@ def rank(self) -> int: "Cannot call rank() for sparse PhlowerTensor" ) - return self._shape_pattern.rank_size + return self._phlower_shape.rank_size def n_vertices(self) -> int: """Returns the number of vertices.""" @@ -235,7 +280,7 @@ def n_vertices(self) -> int: "Cannot call n_vertices() for sparse PhlowerTensor" ) - return self._shape_pattern.n_vertices() + return self._phlower_shape.get_n_vertices() def indices(self) -> torch.Tensor: return self._tensor.indices() @@ -253,21 +298,20 @@ def to_vertexwise(self) -> tuple[PhlowerTensor, str]: resultant_pattern : str Pattern of the resultant shape. Can be used for rearrange. """ - original_pattern = self._shape_pattern.pattern + original_pattern = self._phlower_shape.get_pattern() resultant_pattern = ( - f"({self._shape_pattern.space_pattern}) " - f"({self._shape_pattern.time_series_pattern} " - f"{self._shape_pattern.feature_pattern})" + f"({self._phlower_shape.space_pattern}) " + f"({self._phlower_shape.time_series_pattern} " + f"{self._phlower_shape.get_feature_pattern()})" ) tensor_2d = einops.rearrange( self.to_tensor(), f"{original_pattern} -> {resultant_pattern}" ) return ( - PhlowerTensor( + PhlowerTensor.from_pattern( tensor_2d, dimension_tensor=self.dimension, - is_time_series=False, - is_voxel=False, + pattern=resultant_pattern, ), resultant_pattern, ) @@ -275,17 +319,13 @@ def to_vertexwise(self) -> tuple[PhlowerTensor, str]: def rearrange( self, pattern: str, - is_time_series: bool = False, - is_voxel: bool = False, - **kwargs: dict[str, int], + **axes_length: dict[str, int], ) -> PhlowerTensor: - tensor = self.to_tensor() - rearranged = einops.rearrange(tensor, pattern, **kwargs) - return PhlowerTensor( - rearranged, - dimension_tensor=self.dimension, - is_time_series=is_time_series, - is_voxel=is_voxel, + rearranged = einops.rearrange(self._tensor, pattern, **axes_length) + + to_pattern = pattern.split("->")[-1] + return PhlowerTensor.from_pattern( + rearranged, dimension_tensor=self.dimension, pattern=to_pattern ) def reshape( @@ -317,6 +357,13 @@ def detach(self) -> PhlowerTensor: def backward(self) -> None: self._tensor.backward() + def as_pattern(self, pattern: str) -> PhlowerTensor: + return PhlowerTensor.from_pattern( + tensor=self._tensor, + dimension_tensor=self._dimension_tensor, + pattern=pattern, + ) + @classmethod def __torch_function__( cls, diff --git a/src/phlower/_base/tensors/_shape_operation.py b/src/phlower/_base/tensors/_shape_operation.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/phlower/_base/tensors/_shape_operation.py @@ -0,0 +1 @@ + diff --git a/src/phlower/_base/tensors/_tensor_shape.py b/src/phlower/_base/tensors/_tensor_shape.py index 1d73c6e..47f01d3 100644 --- a/src/phlower/_base/tensors/_tensor_shape.py +++ b/src/phlower/_base/tensors/_tensor_shape.py @@ -1,29 +1,74 @@ -import functools +from __future__ import annotations import numpy as np import torch -class PhlowerTensorShapePattern: - def __init__(self, shape: torch.Size, is_time_series: bool, is_voxel: bool): +class PhlowerShapePattern: + @classmethod + def from_pattern( + cls, shape: torch.Size, pattern: str + ) -> PhlowerShapePattern: + _splited = _split_pattern(pattern) + is_time_series = _check_is_time_series(_splited) + is_voxel = _check_is_voxel(_splited, is_time_series) + return PhlowerShapePattern(shape, is_time_series, is_voxel) + + def __init__( + self, + shape: torch.Size, + is_time_series: bool, + is_voxel: bool, + ): self._shape = shape self._is_time_series = is_time_series self._is_voxel = is_voxel - def pattern_to_ndim(self, drop_last: bool = False) -> dict[str, int]: - chars = self.pattern.split(" ") + def get_pattern_to_size(self, drop_last: bool = False) -> dict[str, int]: + chars = self.get_pattern().split(" ") if drop_last: chars.pop() return {c: self._shape[i] for i, c in enumerate(chars)} - def n_vertices(self) -> int: + def get_n_vertices(self) -> int: start = 1 if self._is_time_series else 0 if self._is_voxel: return np.prod(self._shape[start : start + 3]) return self._shape[start] + def get_pattern( + self, for_einsum: bool = False, drop_last: bool = False + ) -> str: + patterns = [ + self.time_series_pattern, + self.space_pattern, + self.get_feature_pattern( + for_einsum=for_einsum, drop_last=drop_last + ), + ] + + new_pattern = " ".join([p for p in patterns if len(p) != 0]) + + if not for_einsum: + return new_pattern + + return "".join(new_pattern.split()) + + def get_feature_pattern( + self, for_einsum: bool = False, drop_last: bool = False + ) -> str: + start = self.feature_start_dim + if for_einsum: + offset = 1 if drop_last else 0 + return _availale_variables(length=self.rank - offset) + else: + return " ".join([f"a{i}" for i in range(len(self._shape[start:]))]) + + def __str__(self): + return f"ShapePattern: {self.get_pattern()}" + @property def is_time_series(self) -> bool: return self._is_time_series @@ -40,16 +85,6 @@ def shape(self) -> torch.Size: def rank_size(self) -> int: return len(self._shape[self.feature_start_dim : -1]) - @functools.cached_property - def pattern(self) -> str: - patterns = [ - self.time_series_pattern, - self.space_pattern, - self.feature_pattern, - ] - - return " ".join([p for p in patterns if len(p) != 0]) - @property def time_series_pattern(self) -> str: return "t" if self._is_time_series else "" @@ -65,7 +100,86 @@ def feature_start_dim(self) -> int: return offset_time + offset_space - @functools.cached_property - def feature_pattern(self) -> str: - start = self.feature_start_dim - return " ".join([f"a{i}" for i in range(len(self._shape[start:]))]) + +def _check_is_time_series(patterns: list[str]) -> bool: + return _match_to_one_word(patterns[0], "t") + + +def _check_is_voxel(patterns: list[str], is_time_series: bool) -> bool: + offset = 1 if is_time_series else 0 + + if len(patterns) < offset + 3: + return False + + is_x = _match_to_one_word(patterns[offset], "x") + is_y = _match_to_one_word(patterns[offset + 1], "y") + is_z = _match_to_one_word(patterns[offset + 2], "z") + + return is_x and is_y and is_z + + +def _match_to_one_word(target: str, char: str) -> bool: + if len(target) == 1: + return target == char + + if target.startswith("(") and target.endswith(")"): + _collect = "".join(target[1:-1].split()) + return _collect == char + + return False + + +def _split_pattern(pattern: str) -> list[str]: + splited: list[str] = [] + index = 0 + + while index < len(pattern): + if pattern[index] == "(": + p, index = _collect_until_brace_end(pattern, index) + splited.append(p) + continue + + if pattern[index] == " ": + index += 1 + continue + + if pattern[index] == ".": + if pattern[index : index + 3] == "...": + splited.append("...") + index += 3 + continue + raise ValueError(f"Invalid Ellipse found. {pattern}") + + splited.append(pattern[index]) + index += 1 + + return splited + + +def _collect_until_brace_end(pattern: str, start: int) -> tuple[str, int]: + index = start + if pattern[index] != "(": + raise ValueError(f"{pattern[index:]} is not start with ( .") + + count = 0 + while index < len(pattern): + if pattern[index] == "(": + count += 1 + + if pattern[index] == ")": + count -= 1 + if count == 0: + return pattern[start : index + 1], index + 1 + + index += 1 + + raise ValueError(f"brace is not closed correctly. {pattern}") + + +def _availale_variables(length: int, start: int = 0) -> str: + # No f, t, x, y, n and z because they are "reserved" + available_variables = "abcdeghijklmopqrsuvw" + + if length > len(available_variables): + raise ValueError(f"Required length too long: {length}") + return available_variables[start : start + length] diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index 9173036..ebaa0a6 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -1,9 +1,8 @@ import torch -from phlower._base._dimension import PhysicalDimensions from phlower._base.tensors import phlower_tensor -from phlower._base.tensors._dimension_tensor import PhlowerDimensionTensor from phlower._base.tensors._interface import IPhlowerTensor +from phlower._base.tensors._phlower_tensor import PhysicDimensionLikeObject from phlower.utils.exceptions import PhlowerIncompatibleTensorError @@ -69,32 +68,20 @@ def spmm( Resultant tensor. """ h, resultant_pattern = x.to_vertexwise() - restore_pattern = f"{resultant_pattern} -> {x.shape_pattern.pattern}" - restore_ndim = x.shape_pattern.pattern_to_ndim(drop_last=True) + restore_pattern = f"{resultant_pattern} -> {x.shape_pattern.get_pattern()}" + restore_axes_length = x.shape_pattern.get_pattern_to_size(drop_last=True) + for _ in range(repeat): h = torch.sparse.mm(sparse, h) - return h.rearrange( - restore_pattern, - is_time_series=x.is_time_series, - is_voxel=x.is_voxel, - **restore_ndim, - ) + return h.rearrange(restore_pattern, **restore_axes_length) def einsum( equation: str, - *args: list[IPhlowerTensor], - dimension: ( - PhysicalDimensions - | PhlowerDimensionTensor - | torch.Tensor - | dict[str, float] - | list[float] - | tuple[float] - | None - ) = None, - is_time_series: bool = False, - is_voxel: bool = False, + *args: IPhlowerTensor, + dimension: PhysicDimensionLikeObject | None = None, + is_time_series: bool | None = None, + is_voxel: bool | None = None, ) -> IPhlowerTensor: """ Compute einsum for phlower tensors. @@ -123,6 +110,14 @@ def einsum( raise PhlowerIncompatibleTensorError( f"{e}\n" f"{equation}, {[a.shape for a in args]}" ) from e + + is_none_time = is_time_series is None + is_none_voxel = is_voxel is None + + if is_none_time and is_none_voxel: + pattern = equation.split("->")[-1] + return phlower_tensor(ret_tensor, dimension=dimension, pattern=pattern) + return phlower_tensor( ret_tensor, dimension=dimension, @@ -132,8 +127,8 @@ def einsum( def _availale_variables(length: int, start: int = 0) -> str: - # No f, t, x, y, and z because they are "reserved" - available_variables = "abcdeghijklmnopqrsuvw" + # No f, t, x, y, n and z because they are "reserved" + available_variables = "abcdeghijklmopqrsuvw" if length > len(available_variables): raise ValueError(f"Required length too long: {length}") @@ -165,17 +160,12 @@ def contraction( ) ret_is_time_series = x.is_time_series or y.is_time_series - time_x = "t" if x.is_time_series else "" - time_y = "t" if y.is_time_series else "" + time_x = x.shape_pattern.time_series_pattern + time_y = y.shape_pattern.time_series_pattern time_ret = "t" if ret_is_time_series else "" # No need to consider y because they should be compatible - if x.is_voxel: - space = "xyz" - is_voxel = True - else: - space = "x" - is_voxel = False + space = x.shape_pattern.space_pattern diff_rank = x.rank() - y.rank() unresolved = _availale_variables(diff_rank) @@ -190,8 +180,6 @@ def contraction( x, y, dimension=dimension, - is_time_series=ret_is_time_series, - is_voxel=is_voxel, ) @@ -216,17 +204,12 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: y_rank = y.rank() ret_is_time_series = x.is_time_series or y.is_time_series - time_x = "t" if x.is_time_series else "" - time_y = "t" if y.is_time_series else "" + time_x = x.shape_pattern.time_series_pattern + time_y = y.shape_pattern.time_series_pattern time_ret = "t" if ret_is_time_series else "" # No need to consider y because they should be compatible - if x.is_voxel: - space = "xyz" - is_voxel = True - else: - space = "x" - is_voxel = False + space = x.shape_pattern.space_pattern x_vars = _availale_variables(x_rank) y_vars = _availale_variables(y_rank, start=x_rank) @@ -240,14 +223,7 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: else: dimension = x.dimension * y.dimension - return einsum( - equation, - x, - y, - dimension=dimension, - is_time_series=ret_is_time_series, - is_voxel=is_voxel, - ) + return einsum(equation, x, y, dimension=dimension) def tensor_times_scalar( @@ -294,11 +270,8 @@ def apply_orthogonal_group( if rank == 0: return tensor - time = "t" if tensor.is_time_series else "" - space = "xyz" if tensor.is_voxel else "x" - start_dim = 1 - start_dim = start_dim + 1 if tensor.is_time_series else start_dim - start_dim = start_dim + 3 if tensor.is_voxel else start_dim + time = tensor.shape_pattern.time_series_pattern + space = tensor.shape_pattern.space_pattern s = _availale_variables(rank * 2) str_ortho = ",".join(a + b for a, b in zip(s[::2], s[1::2], strict=True)) @@ -307,35 +280,26 @@ def apply_orthogonal_group( equation = f"{str_ortho},{str_tensor}->{str_ret}" args = [orthogonal_matrix] * rank + [tensor] - return einsum( - equation, - *args, - dimension=tensor.dimension, - is_time_series=tensor.is_time_series, - is_voxel=tensor.is_voxel, - ) + return einsum(equation, *args, dimension=tensor.dimension) def spatial_sum(tensor: IPhlowerTensor) -> IPhlowerTensor: """Compute sum over space.""" - time = "t" if tensor.is_time_series else "" - space = "xyz" if tensor.is_voxel else "x" + time = tensor.shape_pattern.time_series_pattern + space = tensor.shape_pattern.space_pattern start_space = 1 if tensor.is_time_series else 0 space_width = 3 if tensor.is_voxel else 1 squeezed = einsum( - f"{time}{space}...->{time}...", - tensor, - dimension=tensor.dimension, - is_time_series=tensor.is_time_series, - is_voxel=tensor.is_voxel, + f"{time}{space}...->{time}...", tensor, dimension=tensor.dimension ) # keepdim for _ in range(space_width): - squeezed._tensor = torch.unsqueeze(squeezed._tensor, start_space) - return squeezed + squeezed = torch.unsqueeze(squeezed, start_space) + + return squeezed.as_pattern(f"{time}{space}...") def spatial_mean(tensor: IPhlowerTensor) -> IPhlowerTensor: diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index 0e28beb..3cb45aa 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -130,11 +130,11 @@ def test__scalar_div_tensor(): c = 3.0 / a ap = PhlowerTensor(a, dims) - cp = 3.0 / ap + cp: PhlowerTensor = 3.0 / ap np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) - assert cp._dimension_tensor == desired_dims + assert cp.dimension == desired_dims def test__tanh(): @@ -267,11 +267,13 @@ def test__to_vertexwise_inverse( ) vertexwise, resultant_pattern = phlower_tensor.to_vertexwise() assert len(vertexwise.shape) == 2 - pattern = f"{resultant_pattern} -> {phlower_tensor.shape_pattern.pattern}" - dict_shape = phlower_tensor.shape_pattern.pattern_to_ndim(drop_last=True) - actual = vertexwise.rearrange( - pattern, is_time_series=is_time_series, is_voxel=is_voxel, **dict_shape + pattern = ( + f"{resultant_pattern} -> {phlower_tensor.shape_pattern.get_pattern()}" ) + dict_shape = phlower_tensor.shape_pattern.get_pattern_to_size( + drop_last=True + ) + actual = vertexwise.rearrange(pattern, **dict_shape) np.testing.assert_almost_equal( actual.to_tensor().numpy(), phlower_tensor.to_tensor().numpy() ) diff --git a/tests/test_base/test_tensors/test_tensor_shape.py b/tests/test_base/test_tensors/test_tensor_shape.py index 05e0809..381b736 100644 --- a/tests/test_base/test_tensors/test_tensor_shape.py +++ b/tests/test_base/test_tensors/test_tensor_shape.py @@ -1,6 +1,6 @@ import pytest import torch -from phlower._base.tensors._tensor_shape import PhlowerTensorShapePattern +from phlower._base.tensors._tensor_shape import PhlowerShapePattern @pytest.mark.parametrize( @@ -10,11 +10,11 @@ def test__detect_not_time_series_vertexwise_pattern( shapes: tuple[int], desired: str ): - shape = PhlowerTensorShapePattern( + shape = PhlowerShapePattern( torch.Size(shapes), is_time_series=False, is_voxel=False ) - assert shape.pattern == desired + assert shape.get_pattern() == desired @pytest.mark.parametrize( @@ -28,11 +28,11 @@ def test__detect_not_time_series_vertexwise_pattern( def test__detect_time_series_vertexwise_pattern( shapes: tuple[int], desired: str ): - shape = PhlowerTensorShapePattern( + shape = PhlowerShapePattern( torch.Size(shapes), is_time_series=True, is_voxel=False ) - assert shape.pattern == desired + assert shape.get_pattern() == desired @pytest.mark.parametrize( @@ -46,11 +46,11 @@ def test__detect_time_series_vertexwise_pattern( def test__detect_not_time_series_voxel_pattern( shapes: tuple[int], desired: str ): - shape = PhlowerTensorShapePattern( + shape = PhlowerShapePattern( torch.Size(shapes), is_time_series=False, is_voxel=True ) - assert shape.pattern == desired + assert shape.get_pattern() == desired @pytest.mark.parametrize( @@ -62,8 +62,8 @@ def test__detect_not_time_series_voxel_pattern( ], ) def test__detect_time_series_voxel_pattern(shapes: tuple[int], desired: str): - shape = PhlowerTensorShapePattern( + shape = PhlowerShapePattern( torch.Size(shapes), is_time_series=True, is_voxel=True ) - assert shape.pattern == desired + assert shape.get_pattern() == desired diff --git a/tests/test_nn/test_core_modules/test_gcn.py b/tests/test_nn/test_core_modules/test_gcn.py index 3fb61b5..d0b7c28 100644 --- a/tests/test_nn/test_core_modules/test_gcn.py +++ b/tests/test_nn/test_core_modules/test_gcn.py @@ -37,7 +37,7 @@ def test__gcn(size: tuple[int], is_time_series: bool): activations=["tanh", "identity"], ) - actual = model(phlower_tensors, field_data=dict_supports) + actual: PhlowerTensor = model(phlower_tensors, field_data=dict_supports) assert actual.shape == size assert actual.is_time_series == is_time_series From 2f647952263e93b22b949e239a15dab9d66c2c48 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Sat, 28 Sep 2024 17:20:27 +0900 Subject: [PATCH 60/89] add tests. fix bugs to handle pattern. --- src/phlower/_base/tensors/_phlower_tensor.py | 6 +- src/phlower/_base/tensors/_tensor_shape.py | 98 ++++++++++++------- src/phlower/nn/_core_modules/_functions.py | 12 +-- .../test_tensors/test_tensor_shape.py | 40 ++++++++ 4 files changed, 113 insertions(+), 43 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index ac79d45..c204c6c 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -132,6 +132,10 @@ def from_pattern( dimension_tensor: PhlowerDimensionTensor | None = None, pattern: str | None = None, ) -> PhlowerTensor: + + if pattern is None: + raise ValueError("pattern must be set when calling from_pattern.") + phlower_shape: PhlowerShapePattern = PhlowerShapePattern.from_pattern( tensor.shape, pattern ) @@ -300,7 +304,7 @@ def to_vertexwise(self) -> tuple[PhlowerTensor, str]: """ original_pattern = self._phlower_shape.get_pattern() resultant_pattern = ( - f"({self._phlower_shape.space_pattern}) " + f"({self._phlower_shape.get_space_pattern()}) " f"({self._phlower_shape.time_series_pattern} " f"{self._phlower_shape.get_feature_pattern()})" ) diff --git a/src/phlower/_base/tensors/_tensor_shape.py b/src/phlower/_base/tensors/_tensor_shape.py index 47f01d3..a942833 100644 --- a/src/phlower/_base/tensors/_tensor_shape.py +++ b/src/phlower/_base/tensors/_tensor_shape.py @@ -3,6 +3,8 @@ import numpy as np import torch +from phlower.utils.exceptions import PhlowerIncompatibleTensorError + class PhlowerShapePattern: @classmethod @@ -10,6 +12,12 @@ def from_pattern( cls, shape: torch.Size, pattern: str ) -> PhlowerShapePattern: _splited = _split_pattern(pattern) + if not _check_shape_and_pattern(shape, _splited): + raise PhlowerIncompatibleTensorError( + "Invalid tensor shape and pattern. " + f"shape: {shape}, pattern: {pattern}" + ) + is_time_series = _check_is_time_series(_splited) is_voxel = _check_is_voxel(_splited, is_time_series) return PhlowerShapePattern(shape, is_time_series, is_voxel) @@ -38,37 +46,40 @@ def get_n_vertices(self) -> int: return np.prod(self._shape[start : start + 3]) return self._shape[start] - def get_pattern( - self, for_einsum: bool = False, drop_last: bool = False - ) -> str: + def get_space_pattern(self, omit_space: bool = False) -> str: + if not self._is_voxel: + return "n" + + if omit_space: + return "xyz" + + return "x y z" + + def get_pattern(self) -> str: patterns = [ self.time_series_pattern, - self.space_pattern, - self.get_feature_pattern( - for_einsum=for_einsum, drop_last=drop_last - ), + self.get_space_pattern(), + self.get_feature_pattern(), ] new_pattern = " ".join([p for p in patterns if len(p) != 0]) + return new_pattern - if not for_einsum: - return new_pattern - - return "".join(new_pattern.split()) - - def get_feature_pattern( - self, for_einsum: bool = False, drop_last: bool = False - ) -> str: + def get_feature_pattern(self) -> str: start = self.feature_start_dim - if for_einsum: - offset = 1 if drop_last else 0 - return _availale_variables(length=self.rank - offset) - else: - return " ".join([f"a{i}" for i in range(len(self._shape[start:]))]) + return " ".join([f"a{i}" for i in range(len(self._shape[start:]))]) def __str__(self): return f"ShapePattern: {self.get_pattern()}" + @property + def start_space_index(self) -> int: + return 1 if self._is_time_series else 0 + + @property + def space_width(self) -> int: + return 3 if self.is_voxel else 1 + @property def is_time_series(self) -> bool: return self._is_time_series @@ -89,10 +100,6 @@ def rank_size(self) -> int: def time_series_pattern(self) -> str: return "t" if self._is_time_series else "" - @property - def space_pattern(self) -> str: - return "x y z" if self._is_voxel else "n" - @property def feature_start_dim(self) -> int: offset_time = 1 if self._is_time_series else 0 @@ -101,6 +108,20 @@ def feature_start_dim(self) -> int: return offset_time + offset_space +def _check_shape_and_pattern(shape: torch.Size, patterns: list[str]) -> bool: + if len(shape) == len(patterns): + return True + + if len(shape) < len(patterns): + return False + + contain_ellipse = np.any([("..." in p) for p in patterns]) + if contain_ellipse: + return True + else: + return False + + def _check_is_time_series(patterns: list[str]) -> bool: return _match_to_one_word(patterns[0], "t") @@ -132,8 +153,9 @@ def _match_to_one_word(target: str, char: str) -> bool: def _split_pattern(pattern: str) -> list[str]: splited: list[str] = [] index = 0 + N_ = len(pattern) - while index < len(pattern): + while index < N_: if pattern[index] == "(": p, index = _collect_until_brace_end(pattern, index) splited.append(p) @@ -150,9 +172,22 @@ def _split_pattern(pattern: str) -> list[str]: continue raise ValueError(f"Invalid Ellipse found. {pattern}") - splited.append(pattern[index]) - index += 1 + if pattern[index].isalpha(): + splited.append(pattern[index]) + index += 1 + continue + + if pattern[index].isnumeric(): + if index == 0: + raise ValueError( + "pattern starting with numerics is invalid. " + f"pattern: {pattern}" + ) + splited[-1] += pattern[index] + index += 1 + continue + raise ValueError(f"Unknown pattern: {pattern=}") return splited @@ -174,12 +209,3 @@ def _collect_until_brace_end(pattern: str, start: int) -> tuple[str, int]: index += 1 raise ValueError(f"brace is not closed correctly. {pattern}") - - -def _availale_variables(length: int, start: int = 0) -> str: - # No f, t, x, y, n and z because they are "reserved" - available_variables = "abcdeghijklmopqrsuvw" - - if length > len(available_variables): - raise ValueError(f"Required length too long: {length}") - return available_variables[start : start + length] diff --git a/src/phlower/nn/_core_modules/_functions.py b/src/phlower/nn/_core_modules/_functions.py index ebaa0a6..c129d7e 100644 --- a/src/phlower/nn/_core_modules/_functions.py +++ b/src/phlower/nn/_core_modules/_functions.py @@ -165,7 +165,7 @@ def contraction( time_ret = "t" if ret_is_time_series else "" # No need to consider y because they should be compatible - space = x.shape_pattern.space_pattern + space = x.shape_pattern.get_space_pattern(omit_space=True) diff_rank = x.rank() - y.rank() unresolved = _availale_variables(diff_rank) @@ -209,7 +209,7 @@ def tensor_product(x: IPhlowerTensor, y: IPhlowerTensor) -> IPhlowerTensor: time_ret = "t" if ret_is_time_series else "" # No need to consider y because they should be compatible - space = x.shape_pattern.space_pattern + space = x.shape_pattern.get_space_pattern(omit_space=True) x_vars = _availale_variables(x_rank) y_vars = _availale_variables(y_rank, start=x_rank) @@ -271,7 +271,7 @@ def apply_orthogonal_group( return tensor time = tensor.shape_pattern.time_series_pattern - space = tensor.shape_pattern.space_pattern + space = tensor.shape_pattern.get_space_pattern(omit_space=True) s = _availale_variables(rank * 2) str_ortho = ",".join(a + b for a, b in zip(s[::2], s[1::2], strict=True)) @@ -287,9 +287,9 @@ def spatial_sum(tensor: IPhlowerTensor) -> IPhlowerTensor: """Compute sum over space.""" time = tensor.shape_pattern.time_series_pattern - space = tensor.shape_pattern.space_pattern - start_space = 1 if tensor.is_time_series else 0 - space_width = 3 if tensor.is_voxel else 1 + space = tensor.shape_pattern.get_space_pattern(omit_space=True) + start_space = tensor.shape_pattern.start_space_index + space_width = tensor.shape_pattern.space_width squeezed = einsum( f"{time}{space}...->{time}...", tensor, dimension=tensor.dimension diff --git a/tests/test_base/test_tensors/test_tensor_shape.py b/tests/test_base/test_tensors/test_tensor_shape.py index 381b736..ecbf55e 100644 --- a/tests/test_base/test_tensors/test_tensor_shape.py +++ b/tests/test_base/test_tensors/test_tensor_shape.py @@ -1,6 +1,46 @@ import pytest import torch from phlower._base.tensors._tensor_shape import PhlowerShapePattern +from phlower.utils.exceptions import PhlowerIncompatibleTensorError + + +@pytest.mark.parametrize( + "shape, pattern, desired_time_series, desired_voxel", + [ + ((3, 4, 5, 3), "t x y z", True, True), + ((3, 4, 5, 3), "n...", False, False), + ((3, 4, 5, 3), "t n ...", True, False), + ((3, 4, 5, 3), "x y z ...", False, True), + ((10, 2, 5, 3), "(n)...", False, False), + ((3, 4, 5, 3), "(t) (x) ( y) (z )", True, True), + ((3, 4, 5), "t (x y z) f", True, False), + ((3, 4, 5, 3), "n a100 a2 f", False, False), + ], +) +def test__detect_flag_when_using_from_pattern( + shape: tuple[int], + pattern: str, + desired_time_series: bool, + desired_voxel: bool, +): + shape_pattern = PhlowerShapePattern.from_pattern(shape, pattern) + + assert shape_pattern.is_time_series == desired_time_series + assert shape_pattern.is_voxel == desired_voxel + + +@pytest.mark.parametrize( + "shape, pattern", + [ + ((3, 4, 5), "(x y z)"), + ((3, 4, 5), "x y "), + ((10, 3, 4, 5, 9), "(x y z) d f"), + ((3, 4, 5), "t (x y z)"), + ], +) +def test__raise_error_for_invalid_pattern(shape: tuple[int], pattern: str): + with pytest.raises(PhlowerIncompatibleTensorError): + _ = PhlowerShapePattern.from_pattern(shape, pattern) @pytest.mark.parametrize( From 1de09f2103e040c7d27ca421ba72557bcb8b25b9 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Sat, 28 Sep 2024 17:28:44 +0900 Subject: [PATCH 61/89] fix lint warnings --- src/phlower/_base/tensors/_phlower_tensor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index c204c6c..ccaaa0f 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -132,7 +132,6 @@ def from_pattern( dimension_tensor: PhlowerDimensionTensor | None = None, pattern: str | None = None, ) -> PhlowerTensor: - if pattern is None: raise ValueError("pattern must be set when calling from_pattern.") From 5574d3c917526e4728f9b1a75b567a17b5fae102 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 30 Sep 2024 10:25:22 +0900 Subject: [PATCH 62/89] treat is_dimensionless as property --- src/phlower/_base/tensors/_dimension_tensor.py | 5 +++-- src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 67be76c..6fba429 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -114,6 +114,7 @@ def to(self, device: str | torch.device, non_blocking: bool = False): def detach(self) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(tensor=self._tensor.detach()) + @property def is_dimensionless(self) -> bool: """Return True if the tensor is dimensionless.""" return torch.sum(torch.abs(self._tensor)) < 1e-5 @@ -329,7 +330,7 @@ def concatenate( @dimension_wrap_implements(torch.tanh) def tanh(tensor: PhlowerDimensionTensor) -> PhlowerDimensionTensor: - if not tensor.is_dimensionless(): + if not tensor.is_dimensionless: raise DimensionIncompatibleError( f"Should be dimensionless to apply tanh but {tensor}" ) @@ -340,7 +341,7 @@ def tanh(tensor: PhlowerDimensionTensor) -> PhlowerDimensionTensor: def leaky_relu( tensor: PhlowerDimensionTensor, *args: Any, **kwargs: Any ) -> PhlowerDimensionTensor: - if not tensor.is_dimensionless(): + if not tensor.is_dimensionless: raise DimensionIncompatibleError( f"Should be dimensionless to apply leaky_relu but {tensor}" ) diff --git a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py index 1be1832..c425a0d 100644 --- a/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py +++ b/src/phlower/nn/_core_modules/_similarity_equivariant_mlp.py @@ -168,7 +168,7 @@ def forward( h = h - mean h = self._mlp(phlower_tensor_collection({"h": h})) - assert h.dimension.is_dimensionless() + assert h.dimension.is_dimensionless if self._centering: linear_mean = self._linear_weight( From 3df6667db52d91f06e5ff16992e4ccfc6faf262a Mon Sep 17 00:00:00 2001 From: sakamoto Date: Fri, 4 Oct 2024 19:22:45 +0900 Subject: [PATCH 63/89] extend model input and output to handle time series flag --- src/phlower/_base/array/__init__.py | 32 +++- src/phlower/_base/array/_interface_wrapper.py | 12 +- .../_base/array/dense/_ndarray_wrapper.py | 26 ++- .../array/sparse/_sparse_array_wrapper.py | 19 +- .../collections/arrays/_arrays_dict.py | 52 ++---- src/phlower/data/_collate_fn.py | 11 +- src/phlower/data/_data_loader.py | 7 +- src/phlower/data/_datasets.py | 123 ++++++++++--- src/phlower/services/trainer/_trainer.py | 14 +- src/phlower/settings/__init__.py | 10 +- src/phlower/settings/_group_settings.py | 21 ++- src/phlower/settings/_model_setting.py | 96 +++++++--- tests/samples/settings/with_share_nn.yml | 20 +- tests/test_base/test_functionals/conftest.py | 43 +++-- tests/test_data/test_data_loader.py | 87 ++++++--- tests/test_data/test_datasets.py | 79 +++++--- .../test_services/test_trainer/data/train.yml | 21 ++- .../test_trainer/data/train_batch_size.yml | 25 ++- tests/test_settings/data/e2e/setting1.yml | 25 ++- tests/test_settings/data/e2e/setting2.yml | 64 +++++++ .../test_settings/data/groups/cycle_error.yml | 100 +++++----- .../data/groups/duplicate_keys_error.yml | 85 +++++---- .../data/groups/key_missing_error.yml | 97 ++++++---- .../groups/not_matched_last_node_error.yml | 74 +++++--- .../data/groups/simple_group_in_group.yml | 172 ++++++++++-------- .../data/groups/simple_module.yml | 77 +++++--- .../simple_module_with_simultion_field.yml | 22 ++- tests/test_settings/test_model_settings.py | 28 +-- .../check_concatenator_nodes.yml | 17 +- .../data/gcn_setting/check_gcn_nodes.yml | 20 +- .../data/mlp_setting/check_mlp_nodes.yml | 17 +- .../share_setting/check_gcn_share_nodes.yml | 20 +- .../share_setting/check_mlp_share_nodes.yml | 20 +- .../data/share_setting/with_share_nn.yml | 20 +- 34 files changed, 1021 insertions(+), 535 deletions(-) create mode 100644 tests/test_settings/data/e2e/setting2.yml diff --git a/src/phlower/_base/array/__init__.py b/src/phlower/_base/array/__init__.py index a3733ec..c7aa349 100644 --- a/src/phlower/_base/array/__init__.py +++ b/src/phlower/_base/array/__init__.py @@ -1,19 +1,43 @@ -# flake8: noqa +from logging import getLogger from typing import get_args +import numpy as np + +from phlower._base import PhysicalDimensions from phlower._base.array import dense, sparse from phlower._base.array._interface_wrapper import IPhlowerArray from phlower.utils.typing import ArrayDataType, DenseArrayType, SparseArrayType +_logger = getLogger(__name__) + -def phlower_array(data: ArrayDataType | IPhlowerArray) -> IPhlowerArray: +def phlower_array( + data: ArrayDataType | IPhlowerArray, + is_time_series: bool = False, + is_voxel: bool = False, + dimensions: PhysicalDimensions | None = None, +) -> IPhlowerArray: if isinstance(data, IPhlowerArray): + _logger.warning( + "phlower_array function is called for PhlowerArray." + "time_series and is_voxel in arguments are ignored." + ) return data if isinstance(data, DenseArrayType): - return dense.NdArrayWrapper(data) + return dense.NdArrayWrapper( + data, + is_time_series=is_time_series, + is_voxel=is_voxel, + dimensions=dimensions, + ) if isinstance(data, get_args(SparseArrayType)): - return sparse.SparseArrayWrapper(data) + if is_time_series or is_voxel: + raise ValueError( + "Sparse Array cannot have time series flag and voxel flag." + ) + + return sparse.SparseArrayWrapper(data, dimensions=dimensions) raise ValueError(f"Unsupported data type: {data.__class__}, {type(data)}") diff --git a/src/phlower/_base/array/_interface_wrapper.py b/src/phlower/_base/array/_interface_wrapper.py index 2d71ceb..93737c1 100644 --- a/src/phlower/_base/array/_interface_wrapper.py +++ b/src/phlower/_base/array/_interface_wrapper.py @@ -10,7 +10,9 @@ class IPhlowerArray(metaclass=abc.ABCMeta): @abc.abstractmethod - def __init__(self, data: ArrayDataType) -> None: ... + def __init__( + self, data: ArrayDataType, dimensions: PhysicalDimensions + ) -> None: ... @property @abc.abstractmethod @@ -20,6 +22,10 @@ def is_time_series(self) -> int | None: ... @abc.abstractmethod def is_sparse(self) -> bool: ... + @property + @abc.abstractmethod + def is_voxel(self) -> bool: ... + @property @abc.abstractmethod def shape(self) -> tuple[int]: ... @@ -71,9 +77,7 @@ def to_phlower_tensor( self, device: str | torch.device | None = None, non_blocking: bool = False, - dimension: PhysicalDimensions | None = None, - is_time_series: bool = False, - is_voxel: bool = False, + disable_dimensions: bool = False, ) -> PhlowerTensor: ... @abc.abstractmethod diff --git a/src/phlower/_base/array/dense/_ndarray_wrapper.py b/src/phlower/_base/array/dense/_ndarray_wrapper.py index bcb160d..889731c 100644 --- a/src/phlower/_base/array/dense/_ndarray_wrapper.py +++ b/src/phlower/_base/array/dense/_ndarray_wrapper.py @@ -10,9 +10,17 @@ class NdArrayWrapper(IPhlowerArray): - def __init__(self, data: np.ndarray, is_time_series: bool = False): + def __init__( + self, + data: np.ndarray, + is_time_series: bool = False, + is_voxel: bool = False, + dimensions: PhysicalDimensions | None = None, + ): self.data = data self._is_time_series = is_time_series + self._is_voxel = is_voxel + self._dimensions = dimensions @property def shape(self) -> tuple[int]: @@ -22,6 +30,10 @@ def shape(self) -> tuple[int]: def is_time_series(self) -> bool: return self._is_time_series + @property + def is_voxel(self) -> bool: + return self._is_voxel + @property def is_sparse(self) -> bool: return False @@ -63,17 +75,13 @@ def to_phlower_tensor( self, device: str | torch.device | None = None, non_blocking: bool = False, - dimension: PhysicalDimensions | None = None, - is_time_series: bool | None = None, - is_voxel: bool = False, + disable_dimensions: bool = False, ) -> PhlowerTensor: - if is_time_series is None: - is_time_series = self.is_time_series _tensor = phlower_tensor( tensor=torch.from_numpy(self.data), - dimension=dimension, - is_time_series=is_time_series, - is_voxel=is_voxel, + dimension=None if disable_dimensions else self._dimensions, + is_time_series=self._is_time_series, + is_voxel=self._is_voxel, ) _tensor.to(device=device, non_blocking=non_blocking) return _tensor diff --git a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py index ffbb88c..dfcffe2 100644 --- a/src/phlower/_base/array/sparse/_sparse_array_wrapper.py +++ b/src/phlower/_base/array/sparse/_sparse_array_wrapper.py @@ -17,8 +17,11 @@ class SparseArrayWrapper(IPhlowerArray): - def __init__(self, arr: SparseArrayType) -> None: + def __init__( + self, arr: SparseArrayType, dimensions: PhysicalDimensions | None = None + ) -> None: self._sparse_data = arr + self._dimensions = dimensions @property def is_sparse(self) -> bool: @@ -28,6 +31,10 @@ def is_sparse(self) -> bool: def is_time_series(self) -> bool: return False + @property + def is_voxel(self) -> bool: + return False + @property def row(self) -> np.ndarray: return self._sparse_data.row @@ -95,9 +102,7 @@ def to_phlower_tensor( self, device: str | torch.device | None = None, non_blocking: bool = False, - dimension: PhysicalDimensions | None = None, - is_time_series: bool = False, - is_voxel: bool = False, + disable_dimensions: bool = False, ) -> PhlowerTensor: sparse_tensor = torch.sparse_coo_tensor( torch.stack( @@ -111,9 +116,9 @@ def to_phlower_tensor( ) _tensor = phlower_tensor( tensor=sparse_tensor, - dimension=dimension, - is_time_series=is_time_series, - is_voxel=is_voxel, + dimension=None if disable_dimensions else self._dimensions, + is_time_series=False, + is_voxel=False, ) _tensor = _tensor.coalesce() _tensor.to(device=device, non_blocking=non_blocking) diff --git a/src/phlower/collections/arrays/_arrays_dict.py b/src/phlower/collections/arrays/_arrays_dict.py index df6cea6..ceb0d87 100644 --- a/src/phlower/collections/arrays/_arrays_dict.py +++ b/src/phlower/collections/arrays/_arrays_dict.py @@ -8,7 +8,6 @@ GraphBatchInfo, IPhlowerArray, PhlowerTensor, - PhysicalDimensions, ) from phlower._base._functionals import to_batch @@ -19,28 +18,15 @@ def __init__(self, name: str, data: list[IPhlowerArray]) -> None: self._data = data assert len(self._data) > 0 - self._is_sparse = self._reduce_is_sparse() - self._is_time_series = self._reduce_is_time_series() + self._is_sparse = self._reduce_flag("is_sparse") + self._is_time_series = self._reduce_flag("is_time_series") + self._is_voxel = self._reduce_flag("is_voxel") - def _reduce_is_sparse(self) -> bool: - _is_sparse = np.unique(np.array([v.is_sparse for v in self._data])) - if len(_is_sparse) != 1: - raise ValueError( - "Sparse array and dense array are mixed " - f"in the same variable name; {self._name}" - ) - return _is_sparse.item() - - def _reduce_is_time_series(self) -> bool: - _is_time_series = np.unique( - np.array([v.is_time_series for v in self._data]) - ) - if len(_is_time_series) != 1: - raise ValueError( - "time-series array and steady array are mixed " - f"in the same variable name; {self._name}" - ) - return _is_time_series.item() + def _reduce_flag(self, attr_name: str) -> bool: + flags = np.unique(np.array([getattr(v, attr_name) for v in self._data])) + if len(flags) != 1: + raise ValueError(f"{attr_name} is not unique in the {self._name}") + return flags.item() def __len__(self) -> int: return len(self._data) @@ -55,19 +41,15 @@ def is_sparse(self) -> bool: def to_batched_tensor( self, - device: str | torch.device | None = None, - non_blocking: bool = False, - dimensions: PhysicalDimensions | None = None, - is_time_series: bool = False, - is_voxel: bool = False, + device: str | torch.device, + non_blocking: bool, + disable_dimensions: bool, ) -> PhlowerTensor: tensors = [ v.to_phlower_tensor( device=device, non_blocking=non_blocking, - dimension=dimensions, - is_time_series=is_time_series, - is_voxel=is_voxel, + disable_dimensions=disable_dimensions, ) for v in self._data ] @@ -105,21 +87,15 @@ def get_names(self) -> Sequence[str]: return self._phlower_sequece_dict.keys() def to_batched_tensor( - self, - device: str, - non_blocking: bool, - dimensions: dict[str, PhysicalDimensions] | None = None, + self, device: str, non_blocking: bool, disable_dimensions: bool = False ) -> tuple[dict[str, PhlowerTensor], dict[str, GraphBatchInfo]]: - if dimensions is None: - dimensions = {} - _batched = [ ( name, arr.to_batched_tensor( device=device, non_blocking=non_blocking, - dimensions=dimensions.get(name), + disable_dimensions=disable_dimensions, ), ) for name, arr in self._phlower_sequece_dict.items() diff --git a/src/phlower/data/_collate_fn.py b/src/phlower/data/_collate_fn.py index 1c8148a..a38c76b 100644 --- a/src/phlower/data/_collate_fn.py +++ b/src/phlower/data/_collate_fn.py @@ -1,6 +1,5 @@ import torch -from phlower._base import PhysicalDimensions from phlower.collections.arrays import SequencedDictArray from phlower.data._lumped_data import LumpedArrayData, LumpedTensorData @@ -10,11 +9,11 @@ def __init__( self, device: str | torch.device, non_blocking: bool = False, - dimensions: dict[str, PhysicalDimensions] | None = None, + disable_dimensions: bool = False, ) -> None: self._device = device self._non_blocking = non_blocking - self._dimensions = dimensions + self._disable_dimensions = disable_dimensions def __call__(self, batch: list[LumpedArrayData]) -> LumpedTensorData: inputs = SequencedDictArray([v.x_data for v in batch]) @@ -25,18 +24,18 @@ def __call__(self, batch: list[LumpedArrayData]) -> LumpedTensorData: inputs_tensors, inputs_batch_info = inputs.to_batched_tensor( device=self._device, non_blocking=self._non_blocking, - dimensions=self._dimensions, + disable_dimensions=self._disable_dimensions, ) outputs_tensors, outputs_batch_info = outputs.to_batched_tensor( device=self._device, non_blocking=self._non_blocking, - dimensions=self._dimensions, + disable_dimensions=self._disable_dimensions, ) field_tensors, field_batch_info = field_data.to_batched_tensor( device=self._device, non_blocking=self._non_blocking, - dimensions=self._dimensions, + disable_dimensions=self._disable_dimensions, ) data_directories = [b.data_directory for b in batch] diff --git a/src/phlower/data/_data_loader.py b/src/phlower/data/_data_loader.py index e8837cf..5140cac 100644 --- a/src/phlower/data/_data_loader.py +++ b/src/phlower/data/_data_loader.py @@ -5,10 +5,9 @@ from torch.utils.data import DataLoader from typing_extensions import Self -from phlower._base import PhysicalDimensions from phlower.data._collate_fn import PhlowerCollateFn from phlower.data._datasets import IPhlowerDataset -from phlower.settings._phlower_setting import ( +from phlower.settings import ( PhlowerPredictorSetting, PhlowerTrainerSetting, ) @@ -45,16 +44,14 @@ def create( self, dataset: IPhlowerDataset, *, - variable_dimensions: dict[str, PhysicalDimensions] | None = None, shuffle: bool = True, disable_dimensions: bool = False, drop_last: bool = False, ) -> DataLoader: - dimensions = None if disable_dimensions else variable_dimensions _collate_fn = PhlowerCollateFn( device=self._device, non_blocking=self._non_blocking, - dimensions=dimensions, + disable_dimensions=disable_dimensions, ) random_generator = torch.Generator(device=self._device) diff --git a/src/phlower/data/_datasets.py b/src/phlower/data/_datasets.py index b62f6e3..1f4a53a 100644 --- a/src/phlower/data/_datasets.py +++ b/src/phlower/data/_datasets.py @@ -1,17 +1,19 @@ import abc import pathlib +import numpy as np from torch.utils.data import Dataset -from phlower._base.array import IPhlowerArray +from phlower._base.array import IPhlowerArray, phlower_array from phlower.data._lumped_data import LumpedArrayData from phlower.io import PhlowerDirectory +from phlower.settings import ModelIOSetting +from phlower.settings._model_setting import _MemberSetting class IPhlowerDataset(metaclass=abc.ABCMeta): @abc.abstractmethod - def __getitem__(self, idx: int) -> LumpedArrayData: - raise NotImplementedError() + def __getitem__(self, idx: int) -> LumpedArrayData: ... class LazyPhlowerDataset(Dataset, IPhlowerDataset): @@ -19,20 +21,22 @@ class LazyPhlowerDataset(Dataset, IPhlowerDataset): def __init__( self, - x_variable_names: list[str], - y_variable_names: list[str] | None, + input_settings: list[ModelIOSetting], + label_settings: list[ModelIOSetting] | None, directories: list[pathlib.Path], *, - field_names: list[str] = None, + field_settings: list[ModelIOSetting] | None = None, allow_no_y_data: bool = False, decrypt_key: bytes | None = None, ): - self._x_variable_names = x_variable_names - if y_variable_names is None: - y_variable_names = [] - self._y_varaible_names = y_variable_names + self._input_settings = input_settings + self._label_settings = ( + label_settings if label_settings is not None else [] + ) self._directories = [PhlowerDirectory(d) for d in directories] - self._field_names = field_names if field_names is not None else [] + self._field_settings = ( + field_settings if field_settings is not None else [] + ) self._allow_no_y_data = allow_no_y_data self._decrypt_key = decrypt_key @@ -42,16 +46,16 @@ def __len__(self): def __getitem__(self, idx: int) -> LumpedArrayData: data_directory = self._directories[idx] - x_data = self._load_data( - data_directory, self._x_variable_names, allow_missing=False + x_data = self._setup_data( + data_directory, self._input_settings, allow_missing=False ) - y_data = self._load_data( + y_data = self._setup_data( data_directory, - self._y_varaible_names, + self._label_settings, allow_missing=self._allow_no_y_data, ) - field_data = self._load_data( - data_directory, self._field_names, allow_missing=False + field_data = self._setup_data( + data_directory, self._field_settings, allow_missing=False ) return LumpedArrayData( x_data=x_data, @@ -60,24 +64,85 @@ def __getitem__(self, idx: int) -> LumpedArrayData: data_directory=data_directory, ) - def _load_data( + def _setup_data( self, data_directory: PhlowerDirectory, - variable_names: list[str], + variable_settings: list[ModelIOSetting], allow_missing: bool, ) -> dict[str, IPhlowerArray]: - variable_files = [ + arrs = [ ( - name, - data_directory.find_variable_file( - name, allow_missing=allow_missing + setting.name, + self._load_phlower_array( + data_directory, setting, allow_missing=allow_missing ), ) - for name in variable_names + for setting in variable_settings ] - return { - name: phlower_file.load(decrypt_key=self._decrypt_key) - for name, phlower_file in variable_files - if phlower_file is not None - } + return {name: arr for name, arr in arrs if arr is not None} + + def _load_phlower_array( + self, + data_directory: PhlowerDirectory, + io_setting: ModelIOSetting, + allow_missing: bool, + ) -> IPhlowerArray | None: + arr = self._load_ndarray_data( + data_directory=data_directory, + members=io_setting.members, + allow_missing=allow_missing, + ) + + if arr is None: + return None + + return phlower_array( + arr, + is_time_series=io_setting.is_time_series, + is_voxel=io_setting.is_voxel, + dimensions=io_setting.physical_dimension, + ) + + def _load_ndarray_data( + self, + data_directory: PhlowerDirectory, + members: list[_MemberSetting], + allow_missing: bool, + ) -> np.ndarray | None: + arrs: list[np.ndarray] = [] + + for member in members: + file_path = data_directory.find_variable_file( + variable_name=member.name, allow_missing=allow_missing + ) + if file_path is None: + continue + + _arr = file_path.load(decrypt_key=self._decrypt_key).to_numpy() + + if member.n_last_dim is None: + arrs.append(_arr) + continue + + if _arr.shape[-1] == member.n_last_dim: + arrs.append(_arr) + continue + + if member.n_last_dim == 1: + _arr = _arr[:, np.newaxis] + arrs.append(_arr) + continue + + raise ValueError( + "Last dimension does not match. " + f"setting: {member.n_last_dim}, actual: {_arr.shape[-1]}" + ) + + if len(arrs) == 0: + return None + + if len(arrs) == 1: + return arrs[0] + + return np.concatenate(arrs) diff --git a/src/phlower/services/trainer/_trainer.py b/src/phlower/services/trainer/_trainer.py index 9588b93..1d69202 100644 --- a/src/phlower/services/trainer/_trainer.py +++ b/src/phlower/services/trainer/_trainer.py @@ -91,27 +91,25 @@ def train( record_io.write_header() train_dataset = LazyPhlowerDataset( - x_variable_names=self._model_setting.network.get_input_keys(), - y_variable_names=self._model_setting.network.get_output_keys(), - field_names=self._model_setting.fields.field_names, + input_settings=self._model_setting.inputs, + label_settings=self._model_setting.labels, + field_settings=self._model_setting.fields, directories=train_directories, ) validation_dataset = LazyPhlowerDataset( - x_variable_names=self._model_setting.network.get_input_keys(), - y_variable_names=self._model_setting.network.get_output_keys(), - field_names=self._model_setting.fields.field_names, + input_settings=self._model_setting.inputs, + label_settings=self._model_setting.labels, + field_settings=self._model_setting.fields, directories=validation_directories, ) builder = DataLoaderBuilder.from_setting(self._trainer_setting) train_loader = builder.create( train_dataset, - variable_dimensions=self._model_setting.variable_dimensions, disable_dimensions=disable_dimensions, ) validation_loader = builder.create( validation_dataset, - variable_dimensions=self._model_setting.variable_dimensions, disable_dimensions=disable_dimensions, ) diff --git a/src/phlower/settings/__init__.py b/src/phlower/settings/__init__.py index bd10428..f209bb2 100644 --- a/src/phlower/settings/__init__.py +++ b/src/phlower/settings/__init__.py @@ -1,5 +1,11 @@ -from phlower.settings._group_settings import GroupModuleSetting, ModuleSetting -from phlower.settings._model_setting import PhlowerModelSetting +from phlower.settings._group_settings import ( + GroupModuleSetting, + ModuleSetting, +) +from phlower.settings._model_setting import ( + ModelIOSetting, + PhlowerModelSetting, +) from phlower.settings._phlower_setting import ( PhlowerPredictorSetting, PhlowerSetting, diff --git a/src/phlower/settings/_group_settings.py b/src/phlower/settings/_group_settings.py index 380e3a5..93e0b74 100644 --- a/src/phlower/settings/_group_settings.py +++ b/src/phlower/settings/_group_settings.py @@ -25,9 +25,10 @@ @dc.dataclass(frozen=True, config=pydantic.ConfigDict(extra="forbid")) -class ModuleIOSetting: +class GroupIOSetting: name: str - n_dim: int + n_last_dim: int + skip_converge: bool = False # HACK need to fix class GroupModuleSetting( @@ -38,12 +39,12 @@ class GroupModuleSetting( name of group """ - inputs: list[ModuleIOSetting] + inputs: list[GroupIOSetting] """ definition of input varaibles """ - outputs: list[ModuleIOSetting] + outputs: list[GroupIOSetting] """ definition of output varaibles """ @@ -120,7 +121,7 @@ def get_output_keys(self) -> list[str]: return [v.name for v in self.outputs] def get_output_info(self) -> dict[str, int]: - return {output.name: output.n_dim for output in self.outputs} + return {output.name: output.n_last_dim for output in self.outputs} def search_module_setting(self, name: str) -> IPhlowerLayerParameters: for module in self.modules: @@ -141,7 +142,7 @@ def resolve( self._check_keys(*resolved_outputs) self._check_nodes(*resolved_outputs) - input_info = {v.name: v.n_dim for v in self.inputs} + input_info = {v.name: v.n_last_dim for v in self.inputs} try: results = _resolve_modules(input_info, self.modules, self) @@ -190,10 +191,10 @@ def _check_nodes(self, *resolved_outputs: dict[str, int]) -> None: key2node.update(output) for input_item in self.inputs: - if key2node[input_item.name] != input_item.n_dim: + if key2node[input_item.name] != input_item.n_last_dim: raise PhlowerModuleNodeDimSizeError( f"n_dim of {input_item.name} in {self.name} " - f"is {input_item.n_dim}. " + f"is {input_item.n_last_dim}. " "It is not consistent with the precedent modules" ) @@ -215,10 +216,10 @@ def _check_last_outputs(self, *resolved_outputs: dict[str, int]) -> None: "Please check last module in this group." ) - if key2node[self_output.name] != self_output.n_dim: + if key2node[self_output.name] != self_output.n_last_dim: raise PhlowerModuleNodeDimSizeError( f"n_dim of {self_output.name} in {self.name} " - f"is {self_output.n_dim}. " + f"is {self_output.n_last_dim}. " "It is not consistent with the precedent modules" ) diff --git a/src/phlower/settings/_model_setting.py b/src/phlower/settings/_model_setting.py index 127c6c3..c25e5c3 100644 --- a/src/phlower/settings/_model_setting.py +++ b/src/phlower/settings/_model_setting.py @@ -1,36 +1,49 @@ from __future__ import annotations +import itertools + import pydantic -from pipe import uniq +from pipe import select, uniq from typing_extensions import Self -from phlower._base import PhysicalDimensionsClass +from phlower._base import PhysicalDimensions, PhysicalDimensionsClass from phlower.settings._group_settings import GroupModuleSetting -class PhlowerFieldSetting(pydantic.BaseModel): - field_names: list[str] = pydantic.Field(default_factory=list, frozen=True) - """ - name of variables in simulation field which are treated as constant - in calculation. - For example, +class _MemberSetting(pydantic.BaseModel): + name: str + n_last_dim: int | None = None + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(extra="forbid", frozen=True) - * support matrix for input graph structure - * boundary conditions +class ModelIOSetting(pydantic.BaseModel): + name: str + is_time_series: bool = False + is_voxel: bool = False + # time_slices: slice | None = None + physical_dimension: PhysicalDimensionsClass | None = None + """ + physical dimension """ - @pydantic.model_validator(mode="after") - def check_duplicate_names(self) -> Self: - unique_names = list(self.field_names | uniq) + members: list[_MemberSetting] = pydantic.Field(default_factory=list) - if len(unique_names) != len(self.field_names): - raise ValueError( - "Duplicate name is found. A varaible name " - "in field setting must be unique." - ) + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(extra="forbid", frozen=True) - return self + @pydantic.model_validator(mode="before") + @classmethod + def _register_if_empty(cls, values: dict) -> dict: + if (values.get("members") is None) or (len(values.get("members")) == 0): + values["members"] = [{"name": values.get("name")}] + + return values + + @property + def n_last_dim(self) -> int: + return sum(v.n_last_dim for v in self.members) class PhlowerModelSetting(pydantic.BaseModel): @@ -41,16 +54,30 @@ class PhlowerModelSetting(pydantic.BaseModel): dictionary which maps variable name to value """ - network: GroupModuleSetting + inputs: list[ModelIOSetting] """ - define structure of neural network + settings for input feature values """ - fields: PhlowerFieldSetting = pydantic.Field( - default_factory=PhlowerFieldSetting + labels: list[ModelIOSetting] = pydantic.Field( + default_factory=list, frozen=True ) """ - settings for fields dependent on your mesh or graph + settings for output feature value + """ + + fields: list[ModelIOSetting] = pydantic.Field( + default_factory=list, frozen=True + ) + """ + name of variables in simulation field which are treated as constant + in calculation. + For example, support matrix for input graph structure. + """ + + network: GroupModuleSetting + """ + define structure of neural network """ # special keyward to forbid extra fields in pydantic @@ -58,6 +85,24 @@ class PhlowerModelSetting(pydantic.BaseModel): frozen=True, extra="forbid", arbitrary_types_allowed=True ) + @pydantic.model_validator(mode="after") + def check_duplicate_names(self) -> Self: + unique_names = list(self.fields | select(lambda x: x.name) | uniq) + + if len(unique_names) != len(self.fields): + raise ValueError( + "Duplicate name is found. A varaible name " + "in field setting must be unique." + ) + + return self + + def get_name_to_dimensions(self) -> dict[str, PhysicalDimensions | None]: + return { + v.name: v.physical_dimension + for v in itertools.chain(self.inputs, self.labels, self.fields) + } + def resolve(self) -> None: """ Resolve network relationship. Following items are checked. @@ -68,4 +113,5 @@ def resolve(self) -> None: * set positive integer value to the value which is defined as -1 in nodes. """ - self.network.resolve(is_first=True) + _inputs = [{v.name: v.n_last_dim} for v in self.inputs] + self.network.resolve(*_inputs) diff --git a/tests/samples/settings/with_share_nn.yml b/tests/samples/settings/with_share_nn.yml index 09f5f6f..21e09a0 100644 --- a/tests/samples/settings/with_share_nn.yml +++ b/tests/samples/settings/with_share_nn.yml @@ -1,17 +1,31 @@ model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP diff --git a/tests/test_base/test_functionals/conftest.py b/tests/test_base/test_functionals/conftest.py index 26138e4..60ab3e4 100644 --- a/tests/test_base/test_functionals/conftest.py +++ b/tests/test_base/test_functionals/conftest.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from phlower._base import PhysicalDimensions from phlower._base.array import phlower_array from phlower._base.tensors._interface import IPhlowerTensor from scipy import sparse as sp @@ -14,23 +15,27 @@ def create_sparse_tensors() -> ( ] ): def _create( - shapes: list[tuple], dimensions: list[dict] = None + shapes: list[tuple], dimensions: list[dict] | None = None ) -> list[IPhlowerTensor]: rng = np.random.default_rng() - arrs = [ - phlower_array( - sp.random(shape[0], shape[1], density=0.1, random_state=rng) - ) - for shape in shapes - ] if dimensions is None: - return [arr.to_phlower_tensor().coalesce() for arr in arrs] - else: return [ - arr.to_phlower_tensor(dimension=dimensions[i]).coalesce() - for i, arr in enumerate(arrs) + phlower_array( + sp.random( + shape[0], shape[1], density=0.1, random_state=rng + ), + ).to_phlower_tensor() + for shape in shapes ] + return [ + phlower_array( + sp.random(shape[0], shape[1], density=0.1, random_state=rng), + dimensions=PhysicalDimensions(dims), + ).to_phlower_tensor() + for shape, dims in zip(shapes, dimensions, strict=True) + ] + return _create @@ -41,15 +46,19 @@ def create_dense_tensors() -> ( ] ): def _create( - shapes: list[tuple], dimensions: list[dict] = None + shapes: list[tuple], dimensions: list[dict] | None = None ) -> list[IPhlowerTensor]: - arrs = [phlower_array(np.random.rand(*shape)) for shape in shapes] if dimensions is None: - return [arr.to_phlower_tensor() for arr in arrs] - else: return [ - arr.to_phlower_tensor(dimension=dimensions[i]) - for i, arr in enumerate(arrs) + phlower_array(np.random.rand(*shape)).to_phlower_tensor() + for shape in shapes ] + return [ + phlower_array( + np.random.rand(*shape), dimensions=dims + ).to_phlower_tensor() + for shape, dims in zip(shapes, dimensions, strict=True) + ] + return _create diff --git a/tests/test_data/test_data_loader.py b/tests/test_data/test_data_loader.py index 7d85631..2ef2553 100644 --- a/tests/test_data/test_data_loader.py +++ b/tests/test_data/test_data_loader.py @@ -10,7 +10,11 @@ LazyPhlowerDataset, LumpedTensorData, ) -from phlower.settings import PhlowerPredictorSetting, PhlowerTrainerSetting +from phlower.settings import ( + ModelIOSetting, + PhlowerPredictorSetting, + PhlowerTrainerSetting, +) from phlower.utils.enums import ModelSelectionType @@ -61,6 +65,21 @@ def test__create_from_predictor_setting(setting: PhlowerPredictorSetting): assert dataloader._num_workers == setting.num_workers +def _to_modelIO_settings( + names: list[tuple[str, int, dict]] | None, +) -> list[ModelIOSetting] | None: + if names is None: + return None + return [ + ModelIOSetting( + name=v, + physical_dimension=dims, + members=[{"name": v, "n_last_dim": n_dim}], + ) + for v, n_dim, dims in names + ] + + @pytest.mark.parametrize("batch_size", [1, 2, 3]) def test__consider_batch_size( batch_size: int, @@ -71,10 +90,14 @@ def test__consider_batch_size( output_base_directory / v for v in ["data0", "data1", "data2"] ] dataset = LazyPhlowerDataset( - x_variable_names=["x0", "x1", "x2"], - y_variable_names=["y0"], + input_settings=_to_modelIO_settings( + [("x0", 1, {}), ("x1", 1, {}), ("x2", 1, {})] + ), + label_settings=_to_modelIO_settings([("y0", 1, {})]), directories=directories, - field_names=["s0", "s1"], + field_settings=_to_modelIO_settings( + [("s0", None, {}), ("s1", None, {})] + ), ) builder = DataLoaderBuilder( @@ -92,16 +115,18 @@ def test__consider_batch_size( @pytest.mark.parametrize( - "dimensions, disable_dimensions, desired", + "x_variables, y_variables, fields, disable_dimensions, desired", [ ( - { - "x0": {"L": 2, "T": -2}, - "x1": {"M": 2}, - "x2": {"I": 1}, - "y0": {"N": -2}, - "s0": {"I": 1}, - }, + [ + ("x0", 1, {"L": 2, "T": -2}), + ("x1", 1, {"M": 2}), + ("x2", 1, {"I": 1}), + ], + [ + ("y0", 1, {"N": -2}), + ], + [("s0", None, {"I": 1})], False, { "x0": PhysicalDimensions({"L": 2, "T": -2}), @@ -114,7 +139,9 @@ def test__consider_batch_size( ], ) def test__consider_dimensions( - dimensions: dict, + x_variables: list, + y_variables: list, + fields: list, disable_dimensions: bool, desired: dict, create_tmp_dataset: None, @@ -124,10 +151,10 @@ def test__consider_dimensions( output_base_directory / v for v in ["data0", "data1", "data2"] ] dataset = LazyPhlowerDataset( - x_variable_names=["x0", "x1", "x2"], - y_variable_names=["y0"], + input_settings=_to_modelIO_settings(x_variables), + label_settings=_to_modelIO_settings(y_variables), directories=directories, - field_names=["s0"], + field_settings=_to_modelIO_settings(fields), ) builder = DataLoaderBuilder( @@ -139,7 +166,6 @@ def test__consider_dimensions( ) dataloader = builder.create( dataset, - variable_dimensions=dimensions, disable_dimensions=disable_dimensions, ) @@ -159,22 +185,24 @@ def test__consider_dimensions( @pytest.mark.parametrize( - "dimensions, disable_dimensions", + "inputs, labels, fields, disable_dimensions", [ ( - { - "x0": {"L": 2, "T": -2}, - "x1": {"M": 2}, - "x2": {"I": 1}, - "y0": {"N": -2}, - "s0": {"I": 1}, - }, + [ + ("x0", 1, {"L": 2, "T": -2}), + ("x1", 1, {"M": 2}), + ("x2", 1, {"I": 1}), + ], + [("y0", 1, {"N": -2})], + [("s0", None, {"I": 1})], True, ) ], ) def test__not_consider_dimensions( - dimensions: dict, + inputs: list, + labels: list, + fields: list, disable_dimensions: bool, create_tmp_dataset: None, output_base_directory: pathlib.Path, @@ -183,10 +211,10 @@ def test__not_consider_dimensions( output_base_directory / v for v in ["data0", "data1", "data2"] ] dataset = LazyPhlowerDataset( - x_variable_names=["x0", "x1", "x2"], - y_variable_names=["y0"], + input_settings=_to_modelIO_settings(inputs), + label_settings=_to_modelIO_settings(labels), directories=directories, - field_names=["s0"], + field_settings=_to_modelIO_settings(fields), ) builder = DataLoaderBuilder( @@ -198,7 +226,6 @@ def test__not_consider_dimensions( ) dataloader = builder.create( dataset, - variable_dimensions=dimensions, disable_dimensions=disable_dimensions, ) diff --git a/tests/test_data/test_datasets.py b/tests/test_data/test_datasets.py index 174eec7..016130f 100644 --- a/tests/test_data/test_datasets.py +++ b/tests/test_data/test_datasets.py @@ -3,9 +3,21 @@ import numpy as np import pytest from phlower.data import LazyPhlowerDataset +from phlower.settings import ModelIOSetting from phlower.utils.typing import ArrayDataType +def _to_modelIO_settings( + names: list[tuple[str, int]] | None, +) -> list[ModelIOSetting] | None: + if names is None: + return None + return [ + ModelIOSetting(name=v, members=[{"name": v, "n_last_dim": n_dim}]) + for v, n_dim in names + ] + + @pytest.mark.parametrize( "directories, desired", [(["data0", "data1"], 2), (["data0", "data1", "data2"], 3)], @@ -18,21 +30,28 @@ def test__lazy_dataset_length( ): directories = [output_base_directory / v for v in directories] dataset = LazyPhlowerDataset( - x_variable_names=["x0", "x1"], - y_variable_names=["y0"], + input_settings=_to_modelIO_settings([("x0", 1), ("x1", 1)]), + label_settings=_to_modelIO_settings([("y0", 1)]), directories=directories, - field_names=["s0"], + field_settings=_to_modelIO_settings([("s0", 1)]), ) assert len(dataset) == desired @pytest.mark.parametrize( - "x_variable_names, y_variable_names, field_names, directory_names", - [(["x0", "x1", "x2"], ["y0"], ["s0", "s1"], ["data0", "data1", "data2"])], + "x_variables, y_variables, field_names, directory_names", + [ + ( + [("x0", 1), ("x1", 1), ("x2", 1)], + [("y0", 1)], + [("s0", None), ("s1", None)], + ["data0", "data1", "data2"], + ) + ], ) def test__lazy_dataset_getitem( - x_variable_names: list[str], - y_variable_names: list[str], + x_variables: list[str], + y_variables: list[str], field_names: list[str], directory_names: list[str], create_tmp_dataset: None, @@ -40,10 +59,10 @@ def test__lazy_dataset_getitem( ): directories = [output_base_directory / v for v in directory_names] dataset = LazyPhlowerDataset( - x_variable_names=x_variable_names, - y_variable_names=y_variable_names, + input_settings=_to_modelIO_settings(x_variables), + label_settings=_to_modelIO_settings(y_variables), directories=directories, - field_names=field_names, + field_settings=_to_modelIO_settings(field_names), ) assert len(dataset) > 1 @@ -54,17 +73,21 @@ def test__lazy_dataset_getitem( assert data_name in desired - for v_name in x_variable_names: + for v_name, _ in x_variables: + desired_shape = desired[data_name][v_name].shape np.testing.assert_array_almost_equal( - item.x_data[v_name].to_numpy(), desired[data_name][v_name] + item.x_data[v_name].to_numpy().reshape(desired_shape), + desired[data_name][v_name], ) - for v_name in y_variable_names: + for v_name, _ in y_variables: + desired_shape = desired[data_name][v_name].shape np.testing.assert_array_almost_equal( - item.y_data[v_name].to_numpy(), desired[data_name][v_name] + item.y_data[v_name].to_numpy().reshape(desired_shape), + desired[data_name][v_name], ) - for v_name in field_names: + for v_name, _ in field_names: np.testing.assert_array_almost_equal( item.field_data[v_name].to_numpy().todense(), desired[data_name][v_name].todense(), @@ -72,15 +95,25 @@ def test__lazy_dataset_getitem( @pytest.mark.parametrize( - "x_variable_names, y_variable_names, field_names, directory_names", + "x_variables, y_variables, field_names, directory_names", [ - (["x0", "x1", "x2"], None, ["s0", "s1"], ["data0", "data1", "data2"]), - (["x0", "x1", "x2"], ["y3"], ["s0", "s1"], ["data0", "data1", "data2"]), + ( + [("x0", 1), ("x1", 1), ("x2", 1)], + None, + [("s0", None), ("s1", None)], + ["data0", "data1", "data2"], + ), + ( + [("x0", 1), ("x1", 1), ("x2", 1)], + [("y3", 1)], + [("s0", None), ("s1", None)], + ["data0", "data1", "data2"], + ), ], ) def test__lazy_dataset_getitem_when_no_ydata( - x_variable_names: list[str], - y_variable_names: list[str], + x_variables: list[str], + y_variables: list[str], field_names: list[str], directory_names: list[str], create_tmp_dataset: None, @@ -88,10 +121,10 @@ def test__lazy_dataset_getitem_when_no_ydata( ): directories = [output_base_directory / v for v in directory_names] dataset = LazyPhlowerDataset( - x_variable_names=x_variable_names, - y_variable_names=y_variable_names, + input_settings=_to_modelIO_settings(x_variables), + label_settings=_to_modelIO_settings(y_variables), directories=directories, - field_names=field_names, + field_settings=_to_modelIO_settings(field_names), allow_no_y_data=True, ) assert len(dataset) > 1 diff --git a/tests/test_services/test_trainer/data/train.yml b/tests/test_services/test_trainer/data/train.yml index 07e2245..aac8353 100644 --- a/tests/test_services/test_trainer/data/train.yml +++ b/tests/test_services/test_trainer/data/train.yml @@ -13,25 +13,28 @@ training: model: - variable_dimensions: - nodal_initial_u: {"L": 1, "T": -1} - nodal_last_u: {"L": 1, "T": -1} - nodal_nadj: {} - + inputs: + - name: nodal_initial_u + physical_dimension: {"L": 1, "T": -1} + + labels: + - name: nodal_last_u + physical_dimension: {"L": 1, "T": -1} + fields: - field_names: - - "nodal_nadj" + - name: "nodal_nadj" + physical_dimension: {} network: nn_type: GROUP name: DEMO inputs: - name: nodal_initial_u - n_dim: 1 + n_last_dim: 1 outputs: - name: nodal_last_u - n_dim: 1 + n_last_dim: 1 modules: - nn_type: MLP diff --git a/tests/test_services/test_trainer/data/train_batch_size.yml b/tests/test_services/test_trainer/data/train_batch_size.yml index c720955..45901d6 100644 --- a/tests/test_services/test_trainer/data/train_batch_size.yml +++ b/tests/test_services/test_trainer/data/train_batch_size.yml @@ -11,25 +11,34 @@ training: nodal_last_u: "mse" model: - variable_dimensions: - nodal_initial_u: {"L": 1, "T": -1} - nodal_last_u: {"L": 1, "T": -1} - nodal_nadj: {} + inputs: + - name: nodal_initial_u + physical_dimension: {"L": 1, "T": -1} + members: + - name: nodal_initial_u + n_last_dim: 1 + labels: + - name: nodal_last_u + physical_dimension: {"L": 1, "T": -1} + members: + - name: nodal_last_u + n_last_dim: 1 + fields: - field_names: - - "nodal_nadj" + - name: nodal_nadj + physical_dimension: {} network: nn_type: GROUP name: DEMO inputs: - name: nodal_initial_u - n_dim: 1 + n_last_dim: 1 outputs: - name: nodal_last_u - n_dim: 1 + n_last_dim: 1 modules: - nn_type: MLP diff --git a/tests/test_settings/data/e2e/setting1.yml b/tests/test_settings/data/e2e/setting1.yml index 3466be9..c122dc7 100644 --- a/tests/test_settings/data/e2e/setting1.yml +++ b/tests/test_settings/data/e2e/setting1.yml @@ -1,19 +1,32 @@ model: - variable_dimensions: - feature0: {"L": 1, "T": -1} - feature1: {"L": 1, "T": -1} + inputs: + - name: feature0 + physical_dimension: {"L": 1, "T": -1} + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + physical_dimension: {"L": 1, "T": -1} + members: + - name: feature0 + n_last_dim: 12 + + fields: + - name: support1 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP diff --git a/tests/test_settings/data/e2e/setting2.yml b/tests/test_settings/data/e2e/setting2.yml new file mode 100644 index 0000000..9d9b7a2 --- /dev/null +++ b/tests/test_settings/data/e2e/setting2.yml @@ -0,0 +1,64 @@ + + +model: + variable_dimensions: + feature0: {"L": 1, "T": -1} + feature1: {"L": 1, "T": -1} + + inputs: + feature0: + is_time_series: true + is_voxel: true + physics_dimension: {"L": 1, "T": -1} + n_last_dim: 1 + members: + - geo1 + - geo2 + + feature1: + is_time_series: true + is_voxel: true + members: + - name: geo1 + n_last_dim: 1 + + fields: + - name: adj_gcn0 + + labels: + - name: feature1 + is_time_series: true + is_voxel: true + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + - name: feature1 + + outputs: + - name: out_feature0 + n_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 100] + activations: ["Identity", "identity"] + + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/data/groups/cycle_error.yml b/tests/test_settings/data/groups/cycle_error.yml index 15082e4..8f7019a 100644 --- a/tests/test_settings/data/groups/cycle_error.yml +++ b/tests/test_settings/data/groups/cycle_error.yml @@ -1,49 +1,65 @@ model: - nn_type: GROUP - name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + members: + - name: feature0 + n_last_dim: 10 + - name: feature1 - n_dim: 12 - - outputs: - - name: feature0 - n_dim: 5 - - modules: - - nn_type: MLP - name: MLP0 - input_keys: - - feature0 - output_key: mlp0 - destinations: - - GCN0 - nn_parameters: - nodes: [-1, 20, 100] - activations: ["relu", "relu"] + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + is_time_series: false + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: feature0 + n_last_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 100] + activations: ["relu", "relu"] - - nn_type: GCN - name: GCN0 - input_keys: - - mlp0 - destinations: - - GCN1 - output_key: out_feature0 - nn_parameters: - nodes: [-1, 5] - activations: ["relu"] - support_name: support1 + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + destinations: + - GCN1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["relu"] + support_name: support1 - - nn_type: GCN - name: GCN1 - input_keys: - - out_feature0 - destinations: - - MLP0 - output_key: feature0 - nn_parameters: - nodes: [-1, 5] - activations: ["identity"] - support_name: support1 + - nn_type: GCN + name: GCN1 + input_keys: + - out_feature0 + destinations: + - MLP0 + output_key: feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/data/groups/duplicate_keys_error.yml b/tests/test_settings/data/groups/duplicate_keys_error.yml index a4051fc..838052d 100644 --- a/tests/test_settings/data/groups/duplicate_keys_error.yml +++ b/tests/test_settings/data/groups/duplicate_keys_error.yml @@ -1,42 +1,59 @@ model: - nn_type: GROUP - name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + members: + - name: feature0 + n_last_dim: 10 - name: feature1 - n_dim: 12 - - outputs: + members: + - name: feature1 + n_last_dim: 12 + + labels: - name: out_feature0 - n_dim: 11 - - modules: - - nn_type: MLP - name: MLP0 - input_keys: - - feature0 - output_key: mlp0 - destinations: - - Concat - nn_parameters: - nodes: [-1, 20, 100] - activations: ["relu", "relu"] + members: + - name: out_feature0 + n_last_dim: 11 + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 11 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - Concat + nn_parameters: + nodes: [-1, 20, 100] + activations: ["relu", "relu"] - - nn_type: MLP - name: MLP1 - input_keys: - - feature1 - output_key: mlp0 - destinations: - - Concat - nn_parameters: - nodes: [-1, 20, 100] - activations: ["relu", "relu"] + - nn_type: MLP + name: MLP1 + input_keys: + - feature1 + output_key: mlp0 + destinations: + - Concat + nn_parameters: + nodes: [-1, 20, 100] + activations: ["relu", "relu"] - - nn_type: Concatenator - name: Concat - input_keys: - - mlp0 - output_key: out_feature0 + - nn_type: Concatenator + name: Concat + input_keys: + - mlp0 + output_key: out_feature0 diff --git a/tests/test_settings/data/groups/key_missing_error.yml b/tests/test_settings/data/groups/key_missing_error.yml index 0e11fbb..46f790b 100644 --- a/tests/test_settings/data/groups/key_missing_error.yml +++ b/tests/test_settings/data/groups/key_missing_error.yml @@ -1,47 +1,64 @@ model: - nn_type: GROUP - name: CYCLE_MODEL inputs: - name: feature0_1 - n_dim: 10 - - name: feature1_1 - n_dim: 12 - - outputs: + members: + - name: feature0_1 + n_last_dim: 10 + - name: feature0_2 + members: + - name: feature0_2 + n_last_dim: 12 + + labels: - name: feature0 - n_dim: 5 - - modules: - - nn_type: MLP - name: MLP0 - input_keys: - - feature0 - output_key: mlp0 - destinations: - - GCN0 - nn_parameters: - nodes: [-1, 20, 100] - activations: ["relu", "relu"] + members: + - name: feature0 + n_last_dim: 5 + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0_1 + n_last_dim: 10 + - name: feature1_1 + n_last_dim: 12 + + outputs: + - name: feature0 + n_last_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 100] + activations: ["relu", "relu"] - - nn_type: GCN - name: GCN0 - input_keys: - - mlp0 - destinations: - - GCN1 - output_key: out_feature0 - nn_parameters: - nodes: [-1, 5] - activations: ["relu"] - support_name: support1 + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + destinations: + - GCN1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["relu"] + support_name: support1 - - nn_type: GCN - name: GCN1 - input_keys: - - out_feature0 - output_key: feature0 - nn_parameters: - nodes: [-1, 5] - activations: ["identity"] - support_name: support1 + - nn_type: GCN + name: GCN1 + input_keys: + - out_feature0 + output_key: feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/data/groups/not_matched_last_node_error.yml b/tests/test_settings/data/groups/not_matched_last_node_error.yml index df1ece0..1e542af 100644 --- a/tests/test_settings/data/groups/not_matched_last_node_error.yml +++ b/tests/test_settings/data/groups/not_matched_last_node_error.yml @@ -1,35 +1,51 @@ model: - nn_type: GROUP - name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + members: + - name: feature0 + n_last_dim: 10 + - name: feature1 - n_dim: 12 - - outputs: - - name: out_feature0 - n_dim: 11 - - modules: - - nn_type: MLP - name: MLP0 - input_keys: - - feature0 - output_key: mlp0 - destinations: - - GCN0 - nn_parameters: - nodes: [-1, 20, 100] - activations: ["identity", "identity"] + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + is_time_series: false + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 11 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 100] + activations: ["identity", "identity"] - - nn_type: GCN - name: GCN0 - input_keys: - - mlp0 - output_key: out_feature0 - nn_parameters: - nodes: [-1, 5] - support_name: support1 - activations: ["identity"] + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + activations: ["identity"] diff --git a/tests/test_settings/data/groups/simple_group_in_group.yml b/tests/test_settings/data/groups/simple_group_in_group.yml index 30af468..3acaefa 100644 --- a/tests/test_settings/data/groups/simple_group_in_group.yml +++ b/tests/test_settings/data/groups/simple_group_in_group.yml @@ -9,89 +9,105 @@ misc: GCN0: 400 model: - nn_type: GROUP - name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + members: + - name: feature0 + n_last_dim: 10 + - name: feature1 - n_dim: 12 - - outputs: - - name: out_feature0 - n_dim: 5 - - modules: - - nn_type: MLP - name: ENCODER0 - input_keys: - - feature0 - output_key: mlp0 - destinations: - - SUB_GROUP - nn_parameters: - nodes: [-1, 20, 100] - activations: ["relu", "relu"] + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + is_time_series: false + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: MLP + name: ENCODER0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - SUB_GROUP + nn_parameters: + nodes: [-1, 20, 100] + activations: ["relu", "relu"] - - nn_type: MLP - name: ENCODER1 - input_keys: - - feature1 - output_key: mlp1 - destinations: - - SUB_GROUP - nn_parameters: - nodes: [-1, 20, 200] - activations: ["relu", "relu"] + - nn_type: MLP + name: ENCODER1 + input_keys: + - feature1 + output_key: mlp1 + destinations: + - SUB_GROUP + nn_parameters: + nodes: [-1, 20, 200] + activations: ["relu", "relu"] - - nn_type: GROUP - name: SUB_GROUP - inputs: - - name: mlp0 - n_dim: 100 - - name: mlp1 - n_dim: 200 - outputs: - - name: mlp2 - n_dim: 400 - destinations: - - GCN0 - modules: - - nn_type: MLP - name: MLP0 - input_keys: - - mlp0 - output_key: mlp0 - destinations: - - CONCAT - nn_parameters: - nodes: [-1, 20, 200] - activations: ["relu", "relu"] + - nn_type: GROUP + name: SUB_GROUP + inputs: + - name: mlp0 + n_last_dim: 100 + - name: mlp1 + n_last_dim: 200 + outputs: + - name: mlp2 + n_last_dim: 400 + destinations: + - GCN0 + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - mlp0 + output_key: mlp0 + destinations: + - CONCAT + nn_parameters: + nodes: [-1, 20, 200] + activations: ["relu", "relu"] - - nn_type: MLP - name: MLP1 - input_keys: - - mlp1 - output_key: mlp1 - destinations: - - CONCAT - nn_parameters: - nodes: [-1, 20, 200] - activations: ["relu", "relu"] + - nn_type: MLP + name: MLP1 + input_keys: + - mlp1 + output_key: mlp1 + destinations: + - CONCAT + nn_parameters: + nodes: [-1, 20, 200] + activations: ["relu", "relu"] - - nn_type: Concatenator - name: CONCAT - input_keys: - - mlp0 - - mlp1 - output_key: mlp2 + - nn_type: Concatenator + name: CONCAT + input_keys: + - mlp0 + - mlp1 + output_key: mlp2 - - nn_type: GCN - name: GCN0 - input_keys: - - mlp2 - output_key: out_feature0 - nn_parameters: - nodes: [-1, 5] - activations: ["identity"] - support_name: support1 + - nn_type: GCN + name: GCN0 + input_keys: + - mlp2 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/data/groups/simple_module.yml b/tests/test_settings/data/groups/simple_module.yml index 991f656..788c0d9 100644 --- a/tests/test_settings/data/groups/simple_module.yml +++ b/tests/test_settings/data/groups/simple_module.yml @@ -6,36 +6,53 @@ misc: model: - nn_type: GROUP - name: CYCLE_MODEL + inputs: - name: feature0 - n_dim: 10 + members: + - name: feature0 + n_last_dim: 10 + - name: feature1 - n_dim: 12 - - outputs: - - name: out_feature0 - n_dim: 5 - - modules: - - nn_type: MLP - name: MLP0 - input_keys: - - feature0 - output_key: mlp0 - destinations: - - GCN0 - nn_parameters: - nodes: [-1, 20, 100] - activations: ["Identity", "identity"] - - - nn_type: GCN - name: GCN0 - input_keys: - - mlp0 - output_key: out_feature0 - nn_parameters: - nodes: [-1, 5] - activations: ["identity"] - support_name: support1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + is_time_series: false + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: MLP + name: MLP0 + input_keys: + - feature0 + output_key: mlp0 + destinations: + - GCN0 + nn_parameters: + nodes: [-1, 20, 100] + activations: ["Identity", "identity"] + + - nn_type: GCN + name: GCN0 + input_keys: + - mlp0 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + activations: ["identity"] + support_name: support1 diff --git a/tests/test_settings/data/groups/simple_module_with_simultion_field.yml b/tests/test_settings/data/groups/simple_module_with_simultion_field.yml index 310ebb3..877248a 100644 --- a/tests/test_settings/data/groups/simple_module_with_simultion_field.yml +++ b/tests/test_settings/data/groups/simple_module_with_simultion_field.yml @@ -5,16 +5,20 @@ misc: GCN0: 100 model: - variable_dimensions: - feature0: - - fields: - supports: - - name: support1 - boundaries: - - name: dirichlet_u - - name: neumann_p + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + fields: + - name: support1 + time_series: false network: nn_type: GROUP diff --git a/tests/test_settings/test_model_settings.py b/tests/test_settings/test_model_settings.py index 7fce089..76af545 100644 --- a/tests/test_settings/test_model_settings.py +++ b/tests/test_settings/test_model_settings.py @@ -2,7 +2,11 @@ import pytest import yaml -from phlower.settings import GroupModuleSetting, ModuleSetting +from phlower.settings import ( + GroupModuleSetting, + ModuleSetting, + PhlowerModelSetting, +) from phlower.utils.exceptions import ( PhlowerModuleCycleError, PhlowerModuleDuplicateKeyError, @@ -41,43 +45,43 @@ def _recursive_check( def test__can_resolve_phlower_networks(file_name: str): data = parse_file(file_name) - setting = GroupModuleSetting(**data["model"]) - setting.resolve(is_first=True) + setting = PhlowerModelSetting(**data["model"]) + setting.resolve() - _recursive_check(setting, data["misc"]["tests"]) + _recursive_check(setting.network, data["misc"]["tests"]) @pytest.mark.parametrize("file_name", ["cycle_error.yml"]) def test__detect_cycle_error(file_name: str): data = parse_file(file_name) - setting = GroupModuleSetting(**data["model"]) + setting = PhlowerModelSetting(**data["model"]) with pytest.raises(PhlowerModuleCycleError): - setting.resolve(is_first=True) + setting.resolve() @pytest.mark.parametrize("file_name", ["not_matched_last_node_error.yml"]) def test__detect_ndim_inconsistency(file_name: str): data = parse_file(file_name) - setting = GroupModuleSetting(**data["model"]) + setting = PhlowerModelSetting(**data["model"]) with pytest.raises(PhlowerModuleNodeDimSizeError): - setting.resolve(is_first=True) + setting.resolve() @pytest.mark.parametrize("file_name", ["duplicate_keys_error.yml"]) def test__detect_duplicate_errors(file_name: str): data = parse_file(file_name) - setting = GroupModuleSetting(**data["model"]) + setting = PhlowerModelSetting(**data["model"]) with pytest.raises(PhlowerModuleDuplicateKeyError): - setting.resolve(is_first=True) + setting.resolve() @pytest.mark.parametrize("file_name", ["key_missing_error.yml"]) def test__detect_key_missing(file_name: str): data = parse_file(file_name) - setting = GroupModuleSetting(**data["model"]) + setting = PhlowerModelSetting(**data["model"]) with pytest.raises(PhlowerModuleKeyError): - setting.resolve(is_first=True) + setting.resolve() diff --git a/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml b/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml index 50c0c45..c466460 100644 --- a/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml +++ b/tests/test_settings/test_module_settings/data/concatenator_setting/check_concatenator_nodes.yml @@ -8,18 +8,29 @@ misc: model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP diff --git a/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml b/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml index d04502d..c817977 100644 --- a/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml +++ b/tests/test_settings/test_module_settings/data/gcn_setting/check_gcn_nodes.yml @@ -6,18 +6,32 @@ misc: GCN2: [5, 5] model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: GCN diff --git a/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml b/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml index 26d5c38..1d9ec33 100644 --- a/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml +++ b/tests/test_settings/test_module_settings/data/mlp_setting/check_mlp_nodes.yml @@ -6,18 +6,29 @@ misc: MLP2: [30, 20, 5] model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + network: nn_type: GROUP name: SAMPLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP diff --git a/tests/test_settings/test_module_settings/data/share_setting/check_gcn_share_nodes.yml b/tests/test_settings/test_module_settings/data/share_setting/check_gcn_share_nodes.yml index a6a201b..46af9d1 100644 --- a/tests/test_settings/test_module_settings/data/share_setting/check_gcn_share_nodes.yml +++ b/tests/test_settings/test_module_settings/data/share_setting/check_gcn_share_nodes.yml @@ -4,18 +4,32 @@ misc: Share0: [5, 5] model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP diff --git a/tests/test_settings/test_module_settings/data/share_setting/check_mlp_share_nodes.yml b/tests/test_settings/test_module_settings/data/share_setting/check_mlp_share_nodes.yml index 064546f..a5508ac 100644 --- a/tests/test_settings/test_module_settings/data/share_setting/check_mlp_share_nodes.yml +++ b/tests/test_settings/test_module_settings/data/share_setting/check_mlp_share_nodes.yml @@ -4,18 +4,32 @@ misc: Share0: [10, 20, 10] model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP diff --git a/tests/test_settings/test_module_settings/data/share_setting/with_share_nn.yml b/tests/test_settings/test_module_settings/data/share_setting/with_share_nn.yml index abf8855..40e4afe 100644 --- a/tests/test_settings/test_module_settings/data/share_setting/with_share_nn.yml +++ b/tests/test_settings/test_module_settings/data/share_setting/with_share_nn.yml @@ -1,17 +1,31 @@ model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + network: nn_type: GROUP name: CYCLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP From 8c52f64aef76187301c08050c2edffcbd19a1405 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 7 Oct 2024 18:21:13 +0900 Subject: [PATCH 64/89] fix e2e test --- src/phlower/services/predictor/_predictor.py | 7 +++---- tests/e2e_tests/data/train.yml | 20 +++++++++++--------- tests/e2e_tests/data/train_batch_size.yml | 20 +++++++++++--------- 3 files changed, 25 insertions(+), 22 deletions(-) diff --git a/src/phlower/services/predictor/_predictor.py b/src/phlower/services/predictor/_predictor.py index 0562c9b..2c075b0 100644 --- a/src/phlower/services/predictor/_predictor.py +++ b/src/phlower/services/predictor/_predictor.py @@ -38,16 +38,15 @@ def predict( disable_dimensions: bool = False, ) -> Iterator[IPhlowerTensorCollections]: dataset = LazyPhlowerDataset( - x_variable_names=self._model_setting.network.get_input_keys(), - y_variable_names=self._model_setting.network.get_output_keys(), - field_names=self._model_setting.fields.field_names, + input_settings=self._model_setting.inputs, + label_settings=self._model_setting.labels, + field_settings=self._model_setting.fields, directories=preprocessed_directories, ) builder = DataLoaderBuilder.from_setting(self._predict_setting) data_loader = builder.create( dataset, - variable_dimensions=self._model_setting.variable_dimensions, disable_dimensions=disable_dimensions, shuffle=False, ) diff --git a/tests/e2e_tests/data/train.yml b/tests/e2e_tests/data/train.yml index 41321d0..edecd09 100644 --- a/tests/e2e_tests/data/train.yml +++ b/tests/e2e_tests/data/train.yml @@ -12,25 +12,27 @@ training: model: - variable_dimensions: - nodal_initial_u: {"L": 1, "T": -1} - nodal_last_u: {"L": 1, "T": -1} - nodal_nadj: {} - + inputs: + - name: nodal_initial_u + physical_dimension: {"L": 1, "T": -1} + labels: + - name: nodal_last_u + physical_dimension: {"L": 1, "T": -1} + fields: - field_names: - - "nodal_nadj" + - name: "nodal_nadj" + physical_dimension: {} network: nn_type: GROUP name: DEMO inputs: - name: nodal_initial_u - n_dim: 1 + n_last_dim: 1 outputs: - name: nodal_last_u - n_dim: 1 + n_last_dim: 1 modules: - nn_type: MLP diff --git a/tests/e2e_tests/data/train_batch_size.yml b/tests/e2e_tests/data/train_batch_size.yml index c720955..83014fe 100644 --- a/tests/e2e_tests/data/train_batch_size.yml +++ b/tests/e2e_tests/data/train_batch_size.yml @@ -11,25 +11,27 @@ training: nodal_last_u: "mse" model: - variable_dimensions: - nodal_initial_u: {"L": 1, "T": -1} - nodal_last_u: {"L": 1, "T": -1} - nodal_nadj: {} - + inputs: + - name: nodal_initial_u + physical_dimension: {"L": 1, "T": -1} + labels: + - name: nodal_last_u + physical_dimension: {"L": 1, "T": -1} + fields: - field_names: - - "nodal_nadj" + - name: "nodal_nadj" + physical_dimension: {} network: nn_type: GROUP name: DEMO inputs: - name: nodal_initial_u - n_dim: 1 + n_last_dim: 1 outputs: - name: nodal_last_u - n_dim: 1 + n_last_dim: 1 modules: - nn_type: MLP From a8f061ab633e4d9cd5276491ca6432709e3381c1 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 7 Oct 2024 21:36:16 +0900 Subject: [PATCH 65/89] fix setting files for tutorials --- .../basic_usages/sample_data/e2e/setting.yml | 19 +++++++++++++------ .../basic_usages/sample_data/model/model.yml | 13 ++++++++++--- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/tutorials/basic_usages/sample_data/e2e/setting.yml b/tutorials/basic_usages/sample_data/e2e/setting.yml index c0acf1a..6603e1e 100644 --- a/tutorials/basic_usages/sample_data/e2e/setting.yml +++ b/tutorials/basic_usages/sample_data/e2e/setting.yml @@ -35,21 +35,28 @@ training: nodal_last_u: "mse" model: - variable_dimensions: - nodal_initial_u: {"L": 1, "T": -1} - nodal_last_u: {"L": 1, "T": -1} - nodal_nadj: {} + inputs: + - name: nodal_initial_u + physical_dimension: {"L": 1, "T": -1} + + labels: + - name: nodal_last_u + physical_dimension: {"L": 1, "T": -1} + + fields: + - name: "nodal_nadj" + physical_dimension: {} network: nn_type: GROUP name: DEMO inputs: - name: nodal_initial_u - n_dim: 1 + n_last_dim: 1 outputs: - name: nodal_last_u - n_dim: 1 + n_last_dim: 1 support_names: ["nodal_nadj"] diff --git a/tutorials/basic_usages/sample_data/model/model.yml b/tutorials/basic_usages/sample_data/model/model.yml index de6a000..454355f 100644 --- a/tutorials/basic_usages/sample_data/model/model.yml +++ b/tutorials/basic_usages/sample_data/model/model.yml @@ -1,15 +1,22 @@ model: + inputs: + - name: feature0 + - name: feature1 + + labels: + - name: out_feature0 + network: nn_type: GROUP name: SAMPLE_MODEL inputs: - name: feature0 - n_dim: 10 + n_last_dim: 10 - name: feature1 - n_dim: 12 + n_last_dim: 12 outputs: - name: out_feature0 - n_dim: 5 + n_last_dim: 5 modules: - nn_type: MLP From 471ceeebec66310a256f7e816c49397024c10e03 Mon Sep 17 00:00:00 2001 From: horiem Date: Tue, 22 Oct 2024 04:02:05 +0900 Subject: [PATCH 66/89] add pytorch like interface --- src/phlower/__init__.py | 2 + .../_base/tensors/_dimension_tensor.py | 12 ++++- src/phlower/_base/tensors/_phlower_tensor.py | 51 ++++++++++++++++--- src/phlower/_base/tensors/_tensor_shape.py | 8 ++- src/phlower/nn/__init__.py | 1 + .../test_tensors/test_phlower_tensor.py | 51 +++++++++++++++++++ .../test_core_modules/test_functions.py | 23 +++++---- 7 files changed, 126 insertions(+), 22 deletions(-) diff --git a/src/phlower/__init__.py b/src/phlower/__init__.py index 1dbbff2..06ad8b1 100644 --- a/src/phlower/__init__.py +++ b/src/phlower/__init__.py @@ -9,4 +9,6 @@ from phlower._fields import ISimulationField from phlower.version import __version__ +from phlower import nn # isort:skip + __all__ = ["__version__"] diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 6fba429..6eed91a 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -108,8 +108,16 @@ def to_dict(self) -> dict[str, float]: for k, v in PhysicalDimensionSymbolType.__members__.items() } - def to(self, device: str | torch.device, non_blocking: bool = False): - self._tensor.to(device, non_blocking=non_blocking) + def to( + self, + device: str | torch.device = None, + non_blocking: bool = False, + dtype: torch.dtype = None, + ) -> PhlowerDimensionTensor: + new_dimension = self._tensor.to( + device, non_blocking=non_blocking, dtype=dtype) + new_dtype = dtype or self.dtype + return PhlowerDimensionTensor(new_dimension, dtype=new_dtype) def detach(self) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(tensor=self._tensor.detach()) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index ccaaa0f..f6936bf 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -57,7 +57,7 @@ def phlower_tensor( def phlower_tensor( - tensor: torch.Tensor | PhlowerTensor, + tensor: list | torch.Tensor | PhlowerTensor, dimension: PhysicDimensionLikeObject | None = None, is_time_series: bool | None = None, is_voxel: bool | None = None, @@ -68,6 +68,9 @@ def phlower_tensor( logger.warning("Input dimension_tensor are ignored.") return tensor + if isinstance(tensor, list): + tensor = torch.tensor(tensor) + dimension_tensor = _resolve_dimension_arg(dimension) if pattern is not None: @@ -130,7 +133,7 @@ def from_pattern( cls, tensor: torch.Tensor, dimension_tensor: PhlowerDimensionTensor | None = None, - pattern: str | None = None, + pattern: str | PhlowerShapePattern | None = None, ) -> PhlowerTensor: if pattern is None: raise ValueError("pattern must be set when calling from_pattern.") @@ -191,12 +194,15 @@ def is_time_series(self) -> bool: def is_voxel(self) -> bool: return self._phlower_shape.is_voxel - def __str__(self) -> str: + def __repr__(self) -> str: return ( f"PhlowerTensor({self._tensor}, " f"Dimension: {self._dimension_tensor})" ) + def __str__(self) -> str: + return self.__repr__() + def __eq__(self, other: PhlowerTensor): return torch.eq(self, other) @@ -219,7 +225,7 @@ def __sub__(self, other: PhlowerTensor): return torch.sub(self, other) def __neg__(self): - return torch.neg(self) + return -1 * self def __add__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.add(self, other) @@ -258,12 +264,23 @@ def to_tensor(self) -> torch.Tensor: def to_numpy(self) -> np.ndarray: return self._tensor.cpu().detach().numpy() + def numpy(self) -> np.ndarray: + return self.to_numpy() + def coalesce(self) -> torch.Tensor: return PhlowerTensor(self._tensor.coalesce(), self._dimension_tensor) def size(self) -> torch.Size: return self._tensor.size() + @property + def dtype(self) -> torch.dtype: + return self._tensor.dtype + + @property + def device(self) -> torch.device: + return self._tensor.device + def numel(self) -> int: return torch.numel(self._tensor) @@ -344,10 +361,21 @@ def reshape( is_voxel=is_voxel, ) - def to(self, device: str, non_blocking: bool = False) -> None: - self._tensor.to(device, non_blocking=non_blocking) + def to( + self, + device: str | torch.device = None, + non_blocking: bool = False, + dtype: torch.dtype = None, + ) -> PhlowerTensor: + new_tensor = self._tensor.to( + device=device, dtype=dtype, non_blocking=non_blocking) if self.has_dimension: - self._dimension_tensor.to(device, non_blocking=non_blocking) + new_dimension = self._dimension_tensor.to( + device=device, dtype=dtype, non_blocking=non_blocking) + else: + new_dimension = None + return PhlowerTensor.from_pattern( + new_tensor, new_dimension, pattern=self.shape_pattern) def detach(self) -> PhlowerTensor: return PhlowerTensor( @@ -367,6 +395,15 @@ def as_pattern(self, pattern: str) -> PhlowerTensor: pattern=pattern, ) + def clone(self) -> PhlowerTensor: + tensor = self._tensor.clone() + return PhlowerTensor( + tensor, + dimension_tensor=self._dimension_tensor, + is_time_series=self.is_time_series, + is_voxel=self.is_voxel, + ) + @classmethod def __torch_function__( cls, diff --git a/src/phlower/_base/tensors/_tensor_shape.py b/src/phlower/_base/tensors/_tensor_shape.py index a942833..7c1d7d9 100644 --- a/src/phlower/_base/tensors/_tensor_shape.py +++ b/src/phlower/_base/tensors/_tensor_shape.py @@ -9,9 +9,13 @@ class PhlowerShapePattern: @classmethod def from_pattern( - cls, shape: torch.Size, pattern: str + cls, shape: torch.Size, pattern: str | PhlowerShapePattern ) -> PhlowerShapePattern: - _splited = _split_pattern(pattern) + if isinstance(pattern, PhlowerShapePattern): + str_pattern = pattern.get_pattern() + else: + str_pattern = pattern + _splited = _split_pattern(str_pattern) if not _check_shape_and_pattern(shape, _splited): raise PhlowerIncompatibleTensorError( "Invalid tensor shape and pattern. " diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 0259afb..8becd39 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -1,3 +1,4 @@ +from phlower.nn._core_modules import _functions as functions from phlower.nn._core_modules._concatenator import Concatenator from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP from phlower.nn._core_modules._gcn import GCN diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index 3cb45aa..b219fb5 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -8,6 +8,37 @@ ) +def test__init(): + list_data = [.1, .2, .3] + pht_list = phlower_tensor(list_data) + pht_torch = phlower_tensor(torch.tensor(list_data)) + np.testing.assert_array_almost_equal( + pht_list.to_numpy(), pht_torch.to_numpy()) + + +def test__to(): + pht = phlower_tensor([.1, .2, .3], dimension={'L': 2, 'T': -1}) + pht16 = pht.to(dtype=torch.float16) + assert pht16._tensor.dtype == torch.float16 + assert pht16.dimension._tensor.dtype == torch.float16 + + +def test__to_numpy(): + pht = phlower_tensor([.1, .2, .3]) + np.testing.assert_array_almost_equal(pht.numpy(), pht.to_numpy()) + + +def test__from_pattern(): + pht = phlower_tensor([.1, .2, .3], dimension={'L': 2, 'T': -1}) + pht_from_pattern = PhlowerTensor.from_pattern( + pht.to_tensor(), pht.dimension, pht.shape_pattern) + + np.testing.assert_array_almost_equal(pht_from_pattern.numpy(), pht.numpy()) + np.testing.assert_array_almost_equal( + pht_from_pattern.dimension._tensor.numpy(), + pht.dimension._tensor.numpy()) + + def test__add(): a = torch.eye(5) b = torch.eye(5) @@ -47,6 +78,17 @@ def test__sub_with_unit(): np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c) +def test__neg_with_unit(): + units = phlower_dimension_tensor({"L": 2, "T": -2}) + a = np.random.rand(3, 10) + c = - a + + ap = PhlowerTensor(torch.tensor(a), units) + cp = - ap + + np.testing.assert_array_almost_equal(cp.numpy(), c) + + @pytest.mark.parametrize( "unit1, unit2", [({"L": 2, "T": -2}, None), (None, {"M": 2, "T": -3})] ) @@ -295,3 +337,12 @@ def test__rearrange( phlower_tensor = PhlowerTensor(torch.rand(*input_shape)) actual = phlower_tensor.rearrange(pattern, **dict_shape) assert actual.shape == desired_shape + + +def test__clone(): + pht = phlower_tensor([.1, .2, .3], dimension={'L': 2, 'T': -1}) + cloned = pht.clone() + pht._tensor[1] = 10. + np.testing.assert_array_almost_equal( + pht.numpy()[[0, 2]], cloned.numpy()[[0, 2]]) + assert pht.numpy()[1] != cloned.numpy()[1] diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index a0c1a9c..f753c7e 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -1,8 +1,8 @@ import numpy as np +import phlower import pytest import torch from phlower import PhlowerTensor, phlower_tensor -from phlower.nn._core_modules import _functions from phlower.utils.exceptions import PhlowerIncompatibleTensorError from scipy import sparse as sp from scipy.stats import ortho_group @@ -33,7 +33,7 @@ def test__spmm(size: tuple[int], is_time_series: bool, repeat: bool): sparse = PhlowerTensor(torch.rand(n, n).to_sparse()) actual_spmm = ( - _functions.spmm(sparse, phlower_tensor, repeat=repeat) + phlower.nn.functions.spmm(sparse, phlower_tensor, repeat=repeat) .to_tensor() .numpy() ) @@ -65,19 +65,20 @@ def assert_correct(actual: np.ndarray, array: np.ndarray): def test_leaky_relu0p5_inverse_leaky_relu0p5(): x = PhlowerTensor(torch.rand(100)) * 4 - 2.0 - y = _functions.inversed_leaky_relu0p5(_functions.leaky_relu0p5(x)) + y = phlower.nn.functions.inversed_leaky_relu0p5( + phlower.nn.functions.leaky_relu0p5(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) def test_tanh_truncated_atanh(): x = PhlowerTensor(torch.rand(100)) * 4 - 2 - y = _functions.truncated_atanh(torch.tanh(x)) + y = phlower.nn.functions.truncated_atanh(torch.tanh(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy(), decimal=6) def test_smooth_leaky_relu_inverse(): x = PhlowerTensor(torch.rand(100)) * 4 - 2 - f = _functions.SmoothLeakyReLU() + f = phlower.nn.functions.SmoothLeakyReLU() y = f.inverse(f(x)) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy(), decimal=5) @@ -126,7 +127,7 @@ def test_contraction_one_argument( is_time_series=is_time_series, is_voxel=is_voxel, ) - actual = _functions.contraction(x) + actual = phlower.nn.functions.contraction(x) desired = torch.einsum(desired_pattern, torch_tensor, torch_tensor).numpy() np.testing.assert_almost_equal(actual.to_tensor().numpy(), desired) @@ -425,7 +426,7 @@ def test_contraction_two_arguments( is_voxel=is_voxel, ) - actual = _functions.contraction(x, y) + actual = phlower.nn.functions.contraction(x, y) desired = torch.einsum(desired_pattern, t_x, t_y).numpy() np.testing.assert_almost_equal(actual.to_tensor().numpy(), desired) @@ -445,7 +446,7 @@ def test_contraction_raises_phlower_incompatible_tensor_error(): x = phlower_tensor(torch.rand(10, 10, 10, 3, 16), is_voxel=True) y = phlower_tensor(torch.rand(10 * 10 * 10, 3, 16), is_voxel=False) with pytest.raises(PhlowerIncompatibleTensorError): - _functions.contraction(x, y) + phlower.nn.functions.contraction(x, y) @pytest.mark.parametrize( @@ -693,7 +694,7 @@ def test_tensor_product( is_voxel=is_voxel, ) - actual = _functions.tensor_product(x, y) + actual = phlower.nn.functions.tensor_product(x, y) desired = torch.einsum(desired_pattern, t_x, t_y).numpy() np.testing.assert_almost_equal(actual.to_numpy(), desired) @@ -756,7 +757,7 @@ def test_apply_orthogonal_group( is_time_series=is_time_series, is_voxel=is_voxel, ) - actual = _functions.apply_orthogonal_group(orthogonal_tensor, x) + actual = phlower.nn.functions.apply_orthogonal_group(orthogonal_tensor, x) if desired_pattern is None: desired = torch_tensor.numpy() @@ -821,7 +822,7 @@ def test_spatial_mean( is_time_series=is_time_series, is_voxel=is_voxel, ) - actual = _functions.spatial_mean(x) + actual = phlower.nn.functions.spatial_mean(x) desired = torch_tensor for dim in mean_dims: From bbf0c2f77a52bdad6d6dd121083fd55a1587598d Mon Sep 17 00:00:00 2001 From: horiem Date: Tue, 22 Oct 2024 04:19:04 +0900 Subject: [PATCH 67/89] update for lint --- .../_base/tensors/_dimension_tensor.py | 11 ++++---- src/phlower/_base/tensors/_phlower_tensor.py | 17 ++++++----- .../test_tensors/test_phlower_tensor.py | 28 +++++++++++-------- .../test_core_modules/test_functions.py | 3 +- 4 files changed, 34 insertions(+), 25 deletions(-) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 6eed91a..45e0d1b 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -109,13 +109,14 @@ def to_dict(self) -> dict[str, float]: } def to( - self, - device: str | torch.device = None, - non_blocking: bool = False, - dtype: torch.dtype = None, + self, + device: str | torch.device = None, + non_blocking: bool = False, + dtype: torch.dtype = None, ) -> PhlowerDimensionTensor: new_dimension = self._tensor.to( - device, non_blocking=non_blocking, dtype=dtype) + device, non_blocking=non_blocking, dtype=dtype + ) new_dtype = dtype or self.dtype return PhlowerDimensionTensor(new_dimension, dtype=new_dtype) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index f6936bf..a19b62e 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -362,20 +362,23 @@ def reshape( ) def to( - self, - device: str | torch.device = None, - non_blocking: bool = False, - dtype: torch.dtype = None, + self, + device: str | torch.device = None, + non_blocking: bool = False, + dtype: torch.dtype = None, ) -> PhlowerTensor: new_tensor = self._tensor.to( - device=device, dtype=dtype, non_blocking=non_blocking) + device=device, dtype=dtype, non_blocking=non_blocking + ) if self.has_dimension: new_dimension = self._dimension_tensor.to( - device=device, dtype=dtype, non_blocking=non_blocking) + device=device, dtype=dtype, non_blocking=non_blocking + ) else: new_dimension = None return PhlowerTensor.from_pattern( - new_tensor, new_dimension, pattern=self.shape_pattern) + new_tensor, new_dimension, pattern=self.shape_pattern + ) def detach(self) -> PhlowerTensor: return PhlowerTensor( diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index b219fb5..a46e3bb 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -9,34 +9,37 @@ def test__init(): - list_data = [.1, .2, .3] + list_data = [0.1, 0.2, 0.3] pht_list = phlower_tensor(list_data) pht_torch = phlower_tensor(torch.tensor(list_data)) np.testing.assert_array_almost_equal( - pht_list.to_numpy(), pht_torch.to_numpy()) + pht_list.to_numpy(), pht_torch.to_numpy() + ) def test__to(): - pht = phlower_tensor([.1, .2, .3], dimension={'L': 2, 'T': -1}) + pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) pht16 = pht.to(dtype=torch.float16) assert pht16._tensor.dtype == torch.float16 assert pht16.dimension._tensor.dtype == torch.float16 def test__to_numpy(): - pht = phlower_tensor([.1, .2, .3]) + pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) np.testing.assert_array_almost_equal(pht.numpy(), pht.to_numpy()) def test__from_pattern(): - pht = phlower_tensor([.1, .2, .3], dimension={'L': 2, 'T': -1}) + pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) pht_from_pattern = PhlowerTensor.from_pattern( - pht.to_tensor(), pht.dimension, pht.shape_pattern) + pht.to_tensor(), pht.dimension, pht.shape_pattern + ) np.testing.assert_array_almost_equal(pht_from_pattern.numpy(), pht.numpy()) np.testing.assert_array_almost_equal( pht_from_pattern.dimension._tensor.numpy(), - pht.dimension._tensor.numpy()) + pht.dimension._tensor.numpy(), + ) def test__add(): @@ -81,10 +84,10 @@ def test__sub_with_unit(): def test__neg_with_unit(): units = phlower_dimension_tensor({"L": 2, "T": -2}) a = np.random.rand(3, 10) - c = - a + c = -a ap = PhlowerTensor(torch.tensor(a), units) - cp = - ap + cp = -ap np.testing.assert_array_almost_equal(cp.numpy(), c) @@ -340,9 +343,10 @@ def test__rearrange( def test__clone(): - pht = phlower_tensor([.1, .2, .3], dimension={'L': 2, 'T': -1}) + pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) cloned = pht.clone() - pht._tensor[1] = 10. + pht._tensor[1] = 10.0 np.testing.assert_array_almost_equal( - pht.numpy()[[0, 2]], cloned.numpy()[[0, 2]]) + pht.numpy()[[0, 2]], cloned.numpy()[[0, 2]] + ) assert pht.numpy()[1] != cloned.numpy()[1] diff --git a/tests/test_nn/test_core_modules/test_functions.py b/tests/test_nn/test_core_modules/test_functions.py index f753c7e..3799348 100644 --- a/tests/test_nn/test_core_modules/test_functions.py +++ b/tests/test_nn/test_core_modules/test_functions.py @@ -66,7 +66,8 @@ def assert_correct(actual: np.ndarray, array: np.ndarray): def test_leaky_relu0p5_inverse_leaky_relu0p5(): x = PhlowerTensor(torch.rand(100)) * 4 - 2.0 y = phlower.nn.functions.inversed_leaky_relu0p5( - phlower.nn.functions.leaky_relu0p5(x)) + phlower.nn.functions.leaky_relu0p5(x) + ) np.testing.assert_almost_equal(x.to_numpy(), y.to_numpy()) From 7f10106b3a24ea7119ac0ce18d82723dace3a90b Mon Sep 17 00:00:00 2001 From: horiem Date: Tue, 22 Oct 2024 04:21:29 +0900 Subject: [PATCH 68/89] add attributes for PhlowerDimensionTensor --- src/phlower/_base/tensors/_dimension_tensor.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 45e0d1b..6f1d3a8 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -123,6 +123,14 @@ def to( def detach(self) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(tensor=self._tensor.detach()) + @property + def dtype(self) -> torch.dtype: + return self._tensor.dtype + + @property + def device(self) -> torch.device: + return self._tensor.device + @property def is_dimensionless(self) -> bool: """Return True if the tensor is dimensionless.""" From 957e2c350df5988f18ab30824078142d2f4035a8 Mon Sep 17 00:00:00 2001 From: horiem Date: Tue, 22 Oct 2024 14:03:41 +0900 Subject: [PATCH 69/89] update for review --- src/phlower/_base/tensors/_phlower_tensor.py | 5 ++- .../test_tensors/test_phlower_tensor.py | 39 +++++++++++++++---- 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index a19b62e..048963c 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -400,9 +400,12 @@ def as_pattern(self, pattern: str) -> PhlowerTensor: def clone(self) -> PhlowerTensor: tensor = self._tensor.clone() + dimension = PhlowerDimensionTensor( + self._dimension_tensor._tensor.clone() + ) return PhlowerTensor( tensor, - dimension_tensor=self._dimension_tensor, + dimension_tensor=dimension, is_time_series=self.is_time_series, is_voxel=self.is_voxel, ) diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index a46e3bb..fdb9ecb 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -8,7 +8,7 @@ ) -def test__init(): +def test__create_same_initialized_object_from_list_and_tensor(): list_data = [0.1, 0.2, 0.3] pht_list = phlower_tensor(list_data) pht_torch = phlower_tensor(torch.tensor(list_data)) @@ -17,14 +17,33 @@ def test__init(): ) -def test__to(): +@pytest.mark.parametrize( + "device", + [ + torch.device("cpu"), + torch.device("meta"), + # TODO: Add CUDA checking + # torch.device('cuda:0'), + # torch.device('cuda:1'), + # torch.device('cuda:2'), + ], +) +@pytest.mark.parametrize( + "dtype", + [torch.float16, torch.float32, torch.float64], +) +def test__to(device: torch.device, dtype: torch.dtype): + print(device, dtype) pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) - pht16 = pht.to(dtype=torch.float16) - assert pht16._tensor.dtype == torch.float16 - assert pht16.dimension._tensor.dtype == torch.float16 + + converted_pht = pht.to(device=device, dtype=dtype) + assert converted_pht.device == device + assert converted_pht.dimension.device == device + assert converted_pht.dtype == dtype + assert converted_pht.dimension.dtype == dtype -def test__to_numpy(): +def test__to_numpy_same_as_numpy(): pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) np.testing.assert_array_almost_equal(pht.numpy(), pht.to_numpy()) @@ -343,10 +362,16 @@ def test__rearrange( def test__clone(): - pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) + original_dimension_dict = {"L": 2, "T": -1} + pht = phlower_tensor([0.1, 0.2, 0.3], dimension=original_dimension_dict) cloned = pht.clone() pht._tensor[1] = 10.0 + pht._dimension_tensor = pht.dimension * pht.dimension np.testing.assert_array_almost_equal( pht.numpy()[[0, 2]], cloned.numpy()[[0, 2]] ) assert pht.numpy()[1] != cloned.numpy()[1] + + for k, v in original_dimension_dict.items(): + assert cloned.dimension.to_dict()[k] == v + assert pht.dimension.to_dict()[k] == 2 * v From 634d99332991de067809e8e3662b199528b15ad6 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 22 Oct 2024 16:31:42 +0900 Subject: [PATCH 70/89] fix import orders --- src/phlower/__init__.py | 3 +-- src/phlower/_fields/_simulation_field.py | 3 +-- src/phlower/nn/_core_modules/_concatenator.py | 2 +- src/phlower/nn/_core_modules/_gcn.py | 2 +- src/phlower/nn/_core_modules/_mlp.py | 2 +- src/phlower/nn/_core_modules/_share.py | 2 +- src/phlower/nn/_group_module.py | 2 +- src/phlower/nn/_interface_module.py | 2 +- src/phlower/nn/_phlower_module_adpter.py | 2 +- 9 files changed, 9 insertions(+), 11 deletions(-) diff --git a/src/phlower/__init__.py b/src/phlower/__init__.py index 06ad8b1..09a03a2 100644 --- a/src/phlower/__init__.py +++ b/src/phlower/__init__.py @@ -1,3 +1,4 @@ +from phlower import nn from phlower._base import ( PhlowerDimensionTensor, PhlowerTensor, @@ -9,6 +10,4 @@ from phlower._fields import ISimulationField from phlower.version import __version__ -from phlower import nn # isort:skip - __all__ = ["__version__"] diff --git a/src/phlower/_fields/_simulation_field.py b/src/phlower/_fields/_simulation_field.py index 0a0a5bb..73cfa84 100644 --- a/src/phlower/_fields/_simulation_field.py +++ b/src/phlower/_fields/_simulation_field.py @@ -1,8 +1,7 @@ import abc from collections.abc import Iterable -from phlower import PhlowerTensor -from phlower._base import GraphBatchInfo +from phlower._base import GraphBatchInfo, PhlowerTensor from phlower.collections.tensors import IPhlowerTensorCollections diff --git a/src/phlower/nn/_core_modules/_concatenator.py b/src/phlower/nn/_core_modules/_concatenator.py index d735f14..d7c5791 100644 --- a/src/phlower/nn/_core_modules/_concatenator.py +++ b/src/phlower/nn/_core_modules/_concatenator.py @@ -3,8 +3,8 @@ import torch from typing_extensions import Self -from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor +from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils from phlower.nn._interface_module import ( diff --git a/src/phlower/nn/_core_modules/_gcn.py b/src/phlower/nn/_core_modules/_gcn.py index f496819..a7a0b77 100644 --- a/src/phlower/nn/_core_modules/_gcn.py +++ b/src/phlower/nn/_core_modules/_gcn.py @@ -2,8 +2,8 @@ import torch -from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor +from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _functions, _utils from phlower.nn._interface_module import ( diff --git a/src/phlower/nn/_core_modules/_mlp.py b/src/phlower/nn/_core_modules/_mlp.py index a5aa43e..e563de8 100644 --- a/src/phlower/nn/_core_modules/_mlp.py +++ b/src/phlower/nn/_core_modules/_mlp.py @@ -3,8 +3,8 @@ import torch from typing_extensions import Self -from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor +from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _utils from phlower.nn._interface_module import ( diff --git a/src/phlower/nn/_core_modules/_share.py b/src/phlower/nn/_core_modules/_share.py index f8fc7ad..3338d5f 100644 --- a/src/phlower/nn/_core_modules/_share.py +++ b/src/phlower/nn/_core_modules/_share.py @@ -2,8 +2,8 @@ import torch -from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor +from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._interface_module import ( IPhlowerCoreModule, diff --git a/src/phlower/nn/_group_module.py b/src/phlower/nn/_group_module.py index 870cc80..42dcac4 100644 --- a/src/phlower/nn/_group_module.py +++ b/src/phlower/nn/_group_module.py @@ -6,7 +6,7 @@ import torch from typing_extensions import Self -from phlower import ISimulationField +from phlower._fields import ISimulationField from phlower.collections.tensors import ( IPhlowerTensorCollections, phlower_tensor_collection, diff --git a/src/phlower/nn/_interface_module.py b/src/phlower/nn/_interface_module.py index 1fc2a48..5e36572 100644 --- a/src/phlower/nn/_interface_module.py +++ b/src/phlower/nn/_interface_module.py @@ -5,8 +5,8 @@ from typing_extensions import Self -from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor +from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.settings._module_settings import IPhlowerLayerParameters diff --git a/src/phlower/nn/_phlower_module_adpter.py b/src/phlower/nn/_phlower_module_adpter.py index 434e895..d81b13a 100644 --- a/src/phlower/nn/_phlower_module_adpter.py +++ b/src/phlower/nn/_phlower_module_adpter.py @@ -4,7 +4,7 @@ import torch -from phlower import ISimulationField +from phlower._fields import ISimulationField from phlower.collections.tensors import ( IPhlowerTensorCollections, phlower_tensor_collection, From 177d8c1509668b4a54b0de03f6e3da89eacfc482 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 22 Oct 2024 17:04:50 +0900 Subject: [PATCH 71/89] add mock tests assuming gpu environment --- .../test_tensors/test_phlower_tensor.py | 56 ++++++++++++++----- 1 file changed, 43 insertions(+), 13 deletions(-) diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index fdb9ecb..5cf6b97 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -1,3 +1,5 @@ +from unittest import mock + import numpy as np import pytest import torch @@ -19,30 +21,58 @@ def test__create_same_initialized_object_from_list_and_tensor(): @pytest.mark.parametrize( "device", - [ - torch.device("cpu"), - torch.device("meta"), - # TODO: Add CUDA checking - # torch.device('cuda:0'), - # torch.device('cuda:1'), - # torch.device('cuda:2'), - ], + [torch.device("cpu"), torch.device("meta"), "cpu", "meta"], ) @pytest.mark.parametrize( "dtype", [torch.float16, torch.float32, torch.float64], ) -def test__to(device: torch.device, dtype: torch.dtype): - print(device, dtype) - pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) +def test__check_dtype_and_device_after_applying_to( + device: torch.device | str, dtype: torch.dtype +): + pht: PhlowerTensor = phlower_tensor( + [0.1, 0.2, 0.3], dimension={"L": 2, "T": -1} + ) converted_pht = pht.to(device=device, dtype=dtype) - assert converted_pht.device == device - assert converted_pht.dimension.device == device + assert converted_pht.device.type == str(device) + assert converted_pht.dimension.device.type == str(device) assert converted_pht.dtype == dtype assert converted_pht.dimension.dtype == dtype +@pytest.mark.parametrize( + "device", + [ + "cpu", + "meta", + "cuda:0", + "cuda:1" + ], +) +@pytest.mark.parametrize( + "dtype", + [torch.float16, torch.float32, torch.float64], +) +def test__pass_arguments_to_torch_function_with_dimension( + device: str, dtype: torch.dtype +): + pht: PhlowerTensor = phlower_tensor( + [0.1, 0.2, 0.3], dimension=None + ) + + with mock.patch.object(torch.Tensor, "to") as mocked: + mocked.return_value = pht._tensor + + _ = pht.to(device=device, dtype=dtype) + + assert mocked.call_count == 1 + + for args in mocked.call_args_list: + assert args.kwargs.get("device") == device + assert args.kwargs.get("dtype") == dtype + + def test__to_numpy_same_as_numpy(): pht = phlower_tensor([0.1, 0.2, 0.3], dimension={"L": 2, "T": -1}) np.testing.assert_array_almost_equal(pht.numpy(), pht.to_numpy()) From 7781b7a7ea9ee4b50e174e2c227d660fb4477f63 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 22 Oct 2024 17:30:58 +0900 Subject: [PATCH 72/89] fix lint warnings --- tests/test_base/test_tensors/test_phlower_tensor.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index 5cf6b97..2bf7590 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -43,12 +43,7 @@ def test__check_dtype_and_device_after_applying_to( @pytest.mark.parametrize( "device", - [ - "cpu", - "meta", - "cuda:0", - "cuda:1" - ], + ["cpu", "meta", "cuda:0", "cuda:1"], ) @pytest.mark.parametrize( "dtype", @@ -57,9 +52,7 @@ def test__check_dtype_and_device_after_applying_to( def test__pass_arguments_to_torch_function_with_dimension( device: str, dtype: torch.dtype ): - pht: PhlowerTensor = phlower_tensor( - [0.1, 0.2, 0.3], dimension=None - ) + pht: PhlowerTensor = phlower_tensor([0.1, 0.2, 0.3], dimension=None) with mock.patch.object(torch.Tensor, "to") as mocked: mocked.return_value = pht._tensor From 3dba2490389f747080eb0d62f59eef83a2ec23e5 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 22 Oct 2024 18:13:11 +0900 Subject: [PATCH 73/89] fix test name --- tests/test_base/test_tensors/test_phlower_tensor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index 2bf7590..3c8fc7d 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -49,9 +49,7 @@ def test__check_dtype_and_device_after_applying_to( "dtype", [torch.float16, torch.float32, torch.float64], ) -def test__pass_arguments_to_torch_function_with_dimension( - device: str, dtype: torch.dtype -): +def test__pass_arguments_to_torch_function(device: str, dtype: torch.dtype): pht: PhlowerTensor = phlower_tensor([0.1, 0.2, 0.3], dimension=None) with mock.patch.object(torch.Tensor, "to") as mocked: From cab31f4eab747f5ef652623f0c40e4d146d2ff51 Mon Sep 17 00:00:00 2001 From: horiem Date: Tue, 22 Oct 2024 23:51:35 +0900 Subject: [PATCH 74/89] add tests for tensors on gpu --- .gitlab-ci.yml | 7 ++ Makefile | 9 +- .../_base/tensors/_dimension_tensor.py | 28 +++++- .../test_tensors/test_phlower_tensor_gpu.py | 98 +++++++++++++++++++ 4 files changed, 136 insertions(+), 6 deletions(-) create mode 100644 tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 5e2d5f6..b2e1fed 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -38,6 +38,13 @@ pytest: - no-gpu - GenuineIntel +gpu_test: + stage: test + script: + - make gpu_test + tags: + - gpu + e2e_test: stage: test script: diff --git a/Makefile b/Makefile index c527355..bda9700 100644 --- a/Makefile +++ b/Makefile @@ -19,14 +19,19 @@ format: .PHONY: test test: - poetry run pytest tests -m "not e2e_test" --cov=src --cov-report term-missing --durations 5 + poetry run pytest tests -m "not e2e_test and not gpu_test" --cov=src --cov-report term-missing --durations 5 -.PHONY: test +.PHONY: e2e_test e2e_test: poetry run pytest tests -m "e2e_test" +.PHONY: gpu_test +gpu_test: + poetry run pytest tests -m "gpu_test" + + .PHONY: lint lint: poetry run python3 -m ruff check --output-format=full diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 6f1d3a8..796254c 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -182,6 +182,16 @@ def add( return PhlowerDimensionTensor(inputs._tensor) + if all( + isinstance(v, float) or v.is_dimensionless + for v in (inputs, other) + ): + for v in (inputs, other): + if isinstance(v, PhlowerDimensionTensor): + device = v.device + break + return zero_dimension_tensor().to(device) + raise DimensionIncompatibleError() @@ -198,6 +208,16 @@ def sub( return PhlowerDimensionTensor(inputs._tensor) + if all( + isinstance(v, float) or v.is_dimensionless + for v in (inputs, other) + ): + for v in (inputs, other): + if isinstance(v, PhlowerDimensionTensor): + device = v.device + break + return zero_dimension_tensor().to(device) + raise DimensionIncompatibleError() @@ -221,12 +241,12 @@ def mul( _input = ( inputs if isinstance(inputs, PhlowerDimensionTensor) - else zero_dimension_tensor() + else zero_dimension_tensor().to(other.device) ) _other = ( other if isinstance(other, PhlowerDimensionTensor) - else zero_dimension_tensor() + else zero_dimension_tensor().to(inputs.device) ) return PhlowerDimensionTensor(_input._tensor + _other._tensor) @@ -238,12 +258,12 @@ def div( _input = ( inputs if isinstance(inputs, PhlowerDimensionTensor) - else zero_dimension_tensor() + else zero_dimension_tensor().to(other.device) ) _other = ( other if isinstance(other, PhlowerDimensionTensor) - else zero_dimension_tensor() + else zero_dimension_tensor().to(inputs.device) ) return PhlowerDimensionTensor(_input._tensor - _other._tensor) diff --git a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py new file mode 100644 index 0000000..d912442 --- /dev/null +++ b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py @@ -0,0 +1,98 @@ +import numpy as np +import pytest +import torch +from phlower import PhlowerTensor, phlower_tensor + + +def generate_random_phlower_tensor_on_gpu( + has_dimension: bool = True) -> PhlowerTensor: + if has_dimension: + dimension = {"L": 2, "T": -1} + else: + dimension = {} + return phlower_tensor(torch.rand(5), dimension=dimension).to( + torch.device('cuda:0')) + + +def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor: + if isinstance(x, PhlowerTensor): + return x.to_tensor() + return x + + +@pytest.mark.parametrize( + "op", [torch.add, torch.sub, torch.mul, torch.div], +) +@pytest.mark.parametrize( + "a, b", + [ + ( + generate_random_phlower_tensor_on_gpu(), + generate_random_phlower_tensor_on_gpu(), + ), + ], +) +def test__op_tensor_tensor_with_unit_on_gpu( + op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float): + c = op(a, b) + + tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) + np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) + assert c.device == torch.device('cuda:0') + assert c.device == torch.device('cuda:0') + + +@pytest.mark.parametrize( + "op", [torch.mul, torch.div], +) +@pytest.mark.parametrize( + "a, b", + [ + ( + 2.3, + generate_random_phlower_tensor_on_gpu(), + ), + ( + generate_random_phlower_tensor_on_gpu(), + 2.3, + ), + ], +) +def test__op_tensor_scalar_with_unit_on_gpu( + op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float): + c = op(a, b) + + tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) + np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) + assert c.device == torch.device('cuda:0') + assert c.device == torch.device('cuda:0') + + +@pytest.mark.parametrize( + "op", [torch.add, torch.sub, torch.mul, torch.div], +) +@pytest.mark.parametrize( + "a, b", + [ + ( + generate_random_phlower_tensor_on_gpu(has_dimension=False), + generate_random_phlower_tensor_on_gpu(has_dimension=False), + ), + ( + 2.3, + generate_random_phlower_tensor_on_gpu(has_dimension=False), + ), + ( + generate_random_phlower_tensor_on_gpu(has_dimension=False), + 2.3, + ), + ], +) +def test__op_tensor_tensor_without_unit_on_gpu( + op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float): + c = op(a, b) + + tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) + np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) + assert c.device == torch.device('cuda:0') + assert c.dimension.device == torch.device('cuda:0') From ca08ac78131ff3d802ae7b0e8fa492a6ec141323 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 00:09:26 +0900 Subject: [PATCH 75/89] update ci image for gpu --- .gitlab-ci.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index b2e1fed..7db8df6 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -39,6 +39,7 @@ pytest: - GenuineIntel gpu_test: + image: nvidia/cuda:12.4.1-devel-ubuntu20.04 stage: test script: - make gpu_test From 474b23b804649917e00d4f2bbf2585e4720228b4 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 00:15:24 +0900 Subject: [PATCH 76/89] update for lint --- .../_base/tensors/_dimension_tensor.py | 10 +---- .../test_tensors/test_phlower_tensor_gpu.py | 45 ++++++++++++------- 2 files changed, 30 insertions(+), 25 deletions(-) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 796254c..5baff94 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -182,10 +182,7 @@ def add( return PhlowerDimensionTensor(inputs._tensor) - if all( - isinstance(v, float) or v.is_dimensionless - for v in (inputs, other) - ): + if all(isinstance(v, float) or v.is_dimensionless for v in (inputs, other)): for v in (inputs, other): if isinstance(v, PhlowerDimensionTensor): device = v.device @@ -208,10 +205,7 @@ def sub( return PhlowerDimensionTensor(inputs._tensor) - if all( - isinstance(v, float) or v.is_dimensionless - for v in (inputs, other) - ): + if all(isinstance(v, float) or v.is_dimensionless for v in (inputs, other)): for v in (inputs, other): if isinstance(v, PhlowerDimensionTensor): device = v.device diff --git a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py index d912442..ca411f6 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py @@ -5,13 +5,15 @@ def generate_random_phlower_tensor_on_gpu( - has_dimension: bool = True) -> PhlowerTensor: + has_dimension: bool = True, +) -> PhlowerTensor: if has_dimension: dimension = {"L": 2, "T": -1} else: dimension = {} return phlower_tensor(torch.rand(5), dimension=dimension).to( - torch.device('cuda:0')) + torch.device("cuda:0") + ) def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor: @@ -20,8 +22,10 @@ def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor: return x +@pytest.mark.gpu_test @pytest.mark.parametrize( - "op", [torch.add, torch.sub, torch.mul, torch.div], + "op", + [torch.add, torch.sub, torch.mul, torch.div], ) @pytest.mark.parametrize( "a, b", @@ -30,20 +34,23 @@ def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor: generate_random_phlower_tensor_on_gpu(), generate_random_phlower_tensor_on_gpu(), ), - ], + ], ) def test__op_tensor_tensor_with_unit_on_gpu( - op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float): + op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float +): c = op(a, b) tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) - assert c.device == torch.device('cuda:0') - assert c.device == torch.device('cuda:0') + assert c.device == torch.device("cuda:0") + assert c.device == torch.device("cuda:0") +@pytest.mark.gpu_test @pytest.mark.parametrize( - "op", [torch.mul, torch.div], + "op", + [torch.mul, torch.div], ) @pytest.mark.parametrize( "a, b", @@ -56,20 +63,23 @@ def test__op_tensor_tensor_with_unit_on_gpu( generate_random_phlower_tensor_on_gpu(), 2.3, ), - ], + ], ) def test__op_tensor_scalar_with_unit_on_gpu( - op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float): + op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float +): c = op(a, b) tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) - assert c.device == torch.device('cuda:0') - assert c.device == torch.device('cuda:0') + assert c.device == torch.device("cuda:0") + assert c.device == torch.device("cuda:0") +@pytest.mark.gpu_test @pytest.mark.parametrize( - "op", [torch.add, torch.sub, torch.mul, torch.div], + "op", + [torch.add, torch.sub, torch.mul, torch.div], ) @pytest.mark.parametrize( "a, b", @@ -86,13 +96,14 @@ def test__op_tensor_scalar_with_unit_on_gpu( generate_random_phlower_tensor_on_gpu(has_dimension=False), 2.3, ), - ], + ], ) def test__op_tensor_tensor_without_unit_on_gpu( - op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float): + op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float +): c = op(a, b) tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) - assert c.device == torch.device('cuda:0') - assert c.dimension.device == torch.device('cuda:0') + assert c.device == torch.device("cuda:0") + assert c.dimension.device == torch.device("cuda:0") From 3c129f40e9fb25f93f24a5f1886ace65b6dd3a50 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 00:19:49 +0900 Subject: [PATCH 77/89] update ci image for gpu --- .gitlab-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 7db8df6..a2e4eac 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -39,7 +39,7 @@ pytest: - GenuineIntel gpu_test: - image: nvidia/cuda:12.4.1-devel-ubuntu20.04 + image: pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime stage: test script: - make gpu_test From 7f2af21884c60542e4aba479b941f6dbf71e3c07 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 00:54:38 +0900 Subject: [PATCH 78/89] update gpu tests to eliminate cuda in fixtures --- pyproject.toml | 1 + .../test_tensors/test_phlower_tensor_gpu.py | 36 ++++++++++--------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6243bc6..ab2162b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ addopts = [ ] markers = [ "e2e_test: marks tests as End-to-End test (deselect with '-m not e2e_test')", + "gpu_test: marks tests as test using GPUs (deselect with '-m not gpu_test')", "need_multicore: marks tests which need multiple cores" ] diff --git a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py index ca411f6..043cb29 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py @@ -4,24 +4,28 @@ from phlower import PhlowerTensor, phlower_tensor -def generate_random_phlower_tensor_on_gpu( +def generate_random_phlower_tensor( has_dimension: bool = True, ) -> PhlowerTensor: if has_dimension: dimension = {"L": 2, "T": -1} else: dimension = {} - return phlower_tensor(torch.rand(5), dimension=dimension).to( - torch.device("cuda:0") - ) + return phlower_tensor(torch.rand(5), dimension=dimension) -def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor: +def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor | float: if isinstance(x, PhlowerTensor): return x.to_tensor() return x +def to_cuda_if_needed(x: PhlowerTensor | float) -> PhlowerTensor | float: + if isinstance(x, PhlowerTensor): + return x.to('cuda:0') + return x + + @pytest.mark.gpu_test @pytest.mark.parametrize( "op", @@ -31,15 +35,15 @@ def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor: "a, b", [ ( - generate_random_phlower_tensor_on_gpu(), - generate_random_phlower_tensor_on_gpu(), + generate_random_phlower_tensor(), + generate_random_phlower_tensor(), ), ], ) def test__op_tensor_tensor_with_unit_on_gpu( op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float ): - c = op(a, b) + c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) @@ -57,10 +61,10 @@ def test__op_tensor_tensor_with_unit_on_gpu( [ ( 2.3, - generate_random_phlower_tensor_on_gpu(), + generate_random_phlower_tensor(), ), ( - generate_random_phlower_tensor_on_gpu(), + generate_random_phlower_tensor(), 2.3, ), ], @@ -68,7 +72,7 @@ def test__op_tensor_tensor_with_unit_on_gpu( def test__op_tensor_scalar_with_unit_on_gpu( op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float ): - c = op(a, b) + c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) @@ -85,15 +89,15 @@ def test__op_tensor_scalar_with_unit_on_gpu( "a, b", [ ( - generate_random_phlower_tensor_on_gpu(has_dimension=False), - generate_random_phlower_tensor_on_gpu(has_dimension=False), + generate_random_phlower_tensor(has_dimension=False), + generate_random_phlower_tensor(has_dimension=False), ), ( 2.3, - generate_random_phlower_tensor_on_gpu(has_dimension=False), + generate_random_phlower_tensor(has_dimension=False), ), ( - generate_random_phlower_tensor_on_gpu(has_dimension=False), + generate_random_phlower_tensor(has_dimension=False), 2.3, ), ], @@ -101,7 +105,7 @@ def test__op_tensor_scalar_with_unit_on_gpu( def test__op_tensor_tensor_without_unit_on_gpu( op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float ): - c = op(a, b) + c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) From 956b52daba1a510f4fcc67c9f1f5a1c3f2c7cd61 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 01:47:11 +0900 Subject: [PATCH 79/89] update for lint --- .../gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py index 043cb29..2144df8 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py @@ -22,7 +22,7 @@ def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor | float: def to_cuda_if_needed(x: PhlowerTensor | float) -> PhlowerTensor | float: if isinstance(x, PhlowerTensor): - return x.to('cuda:0') + return x.to("cuda:0") return x From e1e7ef0310f6cf460303219988ae1302b5cd88ea Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 14:32:32 +0900 Subject: [PATCH 80/89] add tests for PhlowerDimensionTensor on gpu --- .../_base/tensors/_dimension_tensor.py | 31 ++-- .../test_tensors/test_dimensions_gpu.py | 135 ++++++++++++++++++ .../test_tensors/test_phlower_tensor_gpu.py | 4 +- 3 files changed, 158 insertions(+), 12 deletions(-) create mode 100644 tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 5baff94..e1fc918 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -169,6 +169,23 @@ def mean(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(inputs._tensor) +def _determine_float_or_dimensions( + inputs: PhlowerDimensionTensor | float, + other: PhlowerDimensionTensor | float, +) -> tuple[float, PhlowerDimensionTensor]: + if isinstance(inputs, float): + assert isinstance( + other, PhlowerDimensionTensor + ), f"one is float, but the other is {other}" + return inputs, other + + if isinstance(inputs, PhlowerDimensionTensor): + assert isinstance( + other, float + ), f"one is float, but the other is {inputs}" + return other, inputs + + @dimension_wrap_implements(torch.add) def add( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor @@ -183,11 +200,8 @@ def add( return PhlowerDimensionTensor(inputs._tensor) if all(isinstance(v, float) or v.is_dimensionless for v in (inputs, other)): - for v in (inputs, other): - if isinstance(v, PhlowerDimensionTensor): - device = v.device - break - return zero_dimension_tensor().to(device) + _, dim = _determine_float_or_dimensions(inputs, other) + return zero_dimension_tensor().to(dim.device) raise DimensionIncompatibleError() @@ -206,11 +220,8 @@ def sub( return PhlowerDimensionTensor(inputs._tensor) if all(isinstance(v, float) or v.is_dimensionless for v in (inputs, other)): - for v in (inputs, other): - if isinstance(v, PhlowerDimensionTensor): - device = v.device - break - return zero_dimension_tensor().to(device) + _, dim = _determine_float_or_dimensions(inputs, other) + return zero_dimension_tensor().to(dim.device) raise DimensionIncompatibleError() diff --git a/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py new file mode 100644 index 0000000..5cc2257 --- /dev/null +++ b/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py @@ -0,0 +1,135 @@ +import numpy as np +import pytest +import torch +from phlower._base import PhlowerDimensionTensor, phlower_dimension_tensor + + +def generate_phlower_dimension_tensor( + has_dimension: bool = True, inverse: bool = False, +) -> PhlowerDimensionTensor: + if has_dimension: + dimension = {"L": 2, "T": -1} + else: + dimension = {} + if inverse: + dimension = {k: -v for k, v in dimension.items()} + return phlower_dimension_tensor(dimension) + + +def to_tensor_if_needed( + x: PhlowerDimensionTensor | float +) -> torch.Tensor | float: + if isinstance(x, PhlowerDimensionTensor): + return x.to_tensor() + return x + + +def to_cuda_if_needed( + x: PhlowerDimensionTensor | float +) -> PhlowerDimensionTensor | float: + if isinstance(x, PhlowerDimensionTensor): + return x.to("cuda:0") + return x + + +@pytest.mark.gpu_test +@pytest.mark.parametrize( + "op", + [torch.add, torch.sub, torch.mul, torch.div], +) +@pytest.mark.parametrize( + "a, b", + [ + ( + generate_phlower_dimension_tensor(), + generate_phlower_dimension_tensor(), + ), + ], +) +def test__op_tensor_tensor_with_unit_on_gpu( + op: callable, + a: PhlowerDimensionTensor | float, + b: PhlowerDimensionTensor | float, +): + c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) + + cpu_c = op(a, b) + assert c == cpu_c.to("cuda:0") + assert c.device == torch.device("cuda:0") + + +@pytest.mark.gpu_test +@pytest.mark.parametrize( + "a, b", + [ + ( + 2.3, + generate_phlower_dimension_tensor(), + ), + ( + generate_phlower_dimension_tensor(), + 2.3, + ), + ], +) +def test__mul_tensor_scalar_with_unit_on_gpu( + a: PhlowerDimensionTensor | float, + b: PhlowerDimensionTensor | float, +): + c = torch.mul(to_cuda_if_needed(a), to_cuda_if_needed(b)) + + assert c == generate_phlower_dimension_tensor().to("cuda:0") + assert c.device == torch.device("cuda:0") + + +@pytest.mark.gpu_test +def test__div_tensor_scalar_with_unit_on_gpu(): + a = generate_phlower_dimension_tensor() + b = 2.3 + c = torch.div(to_cuda_if_needed(a), to_cuda_if_needed(b)) + + assert c == generate_phlower_dimension_tensor().to("cuda:0") + assert c.device == torch.device("cuda:0") + + +@pytest.mark.gpu_test +def test__div_scalar_tensor_with_unit_on_gpu(): + a = 2.3 + b = generate_phlower_dimension_tensor() + c = torch.div(to_cuda_if_needed(a), to_cuda_if_needed(b)) + + assert c == generate_phlower_dimension_tensor(inverse=True).to("cuda:0") + assert c.device == torch.device("cuda:0") + + +@pytest.mark.gpu_test +@pytest.mark.parametrize( + "op", + [torch.add, torch.sub, torch.mul, torch.div], +) +@pytest.mark.parametrize( + "a, b", + [ + ( + generate_phlower_dimension_tensor(has_dimension=False), + generate_phlower_dimension_tensor(has_dimension=False), + ), + ( + 2.3, + generate_phlower_dimension_tensor(has_dimension=False), + ), + ( + generate_phlower_dimension_tensor(has_dimension=False), + 2.3, + ), + ], +) +def test__op_tensor_tensor_without_unit_on_gpu( + op: callable, + a: PhlowerDimensionTensor | float, + b: PhlowerDimensionTensor | float, +): + c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) + + assert c.device == torch.device("cuda:0") + assert c.is_dimensionless diff --git a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py index 2144df8..2a11e25 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py @@ -48,7 +48,7 @@ def test__op_tensor_tensor_with_unit_on_gpu( tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) assert c.device == torch.device("cuda:0") - assert c.device == torch.device("cuda:0") + assert c.dimension.device == torch.device("cuda:0") @pytest.mark.gpu_test @@ -77,7 +77,7 @@ def test__op_tensor_scalar_with_unit_on_gpu( tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) assert c.device == torch.device("cuda:0") - assert c.device == torch.device("cuda:0") + assert c.dimension.device == torch.device("cuda:0") @pytest.mark.gpu_test From 4e518b83c6e1a4a712b09a643a081a5101abbe04 Mon Sep 17 00:00:00 2001 From: horiem Date: Wed, 23 Oct 2024 18:19:24 +0900 Subject: [PATCH 81/89] lint --- .../test_base/test_tensors/test_dimensions_gpu.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py index 5cc2257..99af196 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py @@ -1,11 +1,11 @@ -import numpy as np import pytest import torch from phlower._base import PhlowerDimensionTensor, phlower_dimension_tensor def generate_phlower_dimension_tensor( - has_dimension: bool = True, inverse: bool = False, + has_dimension: bool = True, + inverse: bool = False, ) -> PhlowerDimensionTensor: if has_dimension: dimension = {"L": 2, "T": -1} @@ -17,7 +17,7 @@ def generate_phlower_dimension_tensor( def to_tensor_if_needed( - x: PhlowerDimensionTensor | float + x: PhlowerDimensionTensor | float, ) -> torch.Tensor | float: if isinstance(x, PhlowerDimensionTensor): return x.to_tensor() @@ -25,7 +25,7 @@ def to_tensor_if_needed( def to_cuda_if_needed( - x: PhlowerDimensionTensor | float + x: PhlowerDimensionTensor | float, ) -> PhlowerDimensionTensor | float: if isinstance(x, PhlowerDimensionTensor): return x.to("cuda:0") From b4cbb756e45b0c0ab84338ea29eb04248849c39a Mon Sep 17 00:00:00 2001 From: sakamoto Date: Thu, 24 Oct 2024 00:15:12 +0900 Subject: [PATCH 82/89] add random tests. restore coverage. --- pyproject.toml | 2 +- .../_base/tensors/_dimension_tensor.py | 20 ++- .../test_tensors/test_dimensions_gpu.py | 150 ++++++------------ .../test_tensors/test_phlower_tensor_gpu.py | 144 +++++++++-------- .../test_base/test_tensors/test_dimensions.py | 24 +++ 5 files changed, 170 insertions(+), 170 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ab2162b..23ac399 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ pytest = "^8.0.2" mypy = "^1.8.0" pytest-cov = "^4.1.0" ruff = "^0.4.10" -hypothesis = "^6.108.2" +hypothesis = {extras = ["numpy"], version = "^6.115.3"} [tool.poetry.group.docs.dependencies] diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index e1fc918..4f5cfc0 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -55,7 +55,7 @@ def from_list( "length of values is not equal to the number of " f"registered dimension types. input: {len(values)}" ) - _tensor = torch.tensor(values, dtype=torch.int64).reshape( + _tensor = torch.tensor(values, dtype=torch.float32).reshape( len(PhysicalDimensionSymbolType), 1 ) return PhlowerDimensionTensor(_tensor) @@ -77,9 +77,15 @@ def __init__( def __add__(self, __value: object): return torch.add(self, __value) + def __radd__(self, __value: object): + return torch.add(self, __value) + def __mul__(self, __value: object): return torch.mul(self, __value) + def __rmul__(self, __value: object): + return torch.mul(self, __value) + def __eq__(self, other: object) -> bool: if not isinstance(other, PhlowerDimensionTensor): return NotImplemented @@ -185,6 +191,8 @@ def _determine_float_or_dimensions( ), f"one is float, but the other is {inputs}" return other, inputs + raise ValueError(f"Unexpected situation. inputs: {inputs}, other: {other}") + @dimension_wrap_implements(torch.add) def add( @@ -199,7 +207,10 @@ def add( return PhlowerDimensionTensor(inputs._tensor) - if all(isinstance(v, float) or v.is_dimensionless for v in (inputs, other)): + if all( + isinstance(v, (int | float)) or v.is_dimensionless + for v in (inputs, other) + ): _, dim = _determine_float_or_dimensions(inputs, other) return zero_dimension_tensor().to(dim.device) @@ -219,7 +230,10 @@ def sub( return PhlowerDimensionTensor(inputs._tensor) - if all(isinstance(v, float) or v.is_dimensionless for v in (inputs, other)): + if all( + isinstance(v, (int | float)) or v.is_dimensionless + for v in (inputs, other) + ): _, dim = _determine_float_or_dimensions(inputs, other) return zero_dimension_tensor().to(dim.device) diff --git a/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py index 99af196..ef0c397 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_dimensions_gpu.py @@ -1,35 +1,26 @@ +from collections.abc import Callable + import pytest import torch -from phlower._base import PhlowerDimensionTensor, phlower_dimension_tensor +from hypothesis import given, settings +from hypothesis import strategies as st +from phlower._base import PhlowerDimensionTensor +from phlower.utils.enums import PhysicalDimensionSymbolType -def generate_phlower_dimension_tensor( - has_dimension: bool = True, - inverse: bool = False, +@st.composite +def random_phlower_dimension_tensor( + draw: Callable, ) -> PhlowerDimensionTensor: - if has_dimension: - dimension = {"L": 2, "T": -1} - else: - dimension = {} - if inverse: - dimension = {k: -v for k, v in dimension.items()} - return phlower_dimension_tensor(dimension) - - -def to_tensor_if_needed( - x: PhlowerDimensionTensor | float, -) -> torch.Tensor | float: - if isinstance(x, PhlowerDimensionTensor): - return x.to_tensor() - return x + dimensions = draw( + st.lists( + elements=st.floats(allow_nan=False, allow_infinity=False), + min_size=len(PhysicalDimensionSymbolType), + max_size=len(PhysicalDimensionSymbolType), + ) + ) - -def to_cuda_if_needed( - x: PhlowerDimensionTensor | float, -) -> PhlowerDimensionTensor | float: - if isinstance(x, PhlowerDimensionTensor): - return x.to("cuda:0") - return x + return PhlowerDimensionTensor.from_list(dimensions).to(device="cuda:0") @pytest.mark.gpu_test @@ -37,68 +28,46 @@ def to_cuda_if_needed( "op", [torch.add, torch.sub, torch.mul, torch.div], ) -@pytest.mark.parametrize( - "a, b", - [ - ( - generate_phlower_dimension_tensor(), - generate_phlower_dimension_tensor(), - ), - ], +@given( + st.lists( + elements=st.floats(width=32, allow_nan=False, allow_infinity=False), + min_size=len(PhysicalDimensionSymbolType), + max_size=len(PhysicalDimensionSymbolType), + ) ) -def test__op_tensor_tensor_with_unit_on_gpu( - op: callable, - a: PhlowerDimensionTensor | float, - b: PhlowerDimensionTensor | float, +@settings(deadline=None) +def test__op_dimension_tensor_with_same_values_on_gpu( + op: Callable[ + [PhlowerDimensionTensor, PhlowerDimensionTensor], PhlowerDimensionTensor + ], + dimensions: list[float], ): - c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) + a = PhlowerDimensionTensor.from_list(dimensions).to("cuda:0") + b = PhlowerDimensionTensor.from_list(dimensions).to("cuda:0") + c = op(a, b) - cpu_c = op(a, b) + cpu_c = op(a.to("cpu"), b.to("cpu")) assert c == cpu_c.to("cuda:0") assert c.device == torch.device("cuda:0") @pytest.mark.gpu_test @pytest.mark.parametrize( - "a, b", - [ - ( - 2.3, - generate_phlower_dimension_tensor(), - ), - ( - generate_phlower_dimension_tensor(), - 2.3, - ), - ], + "op", + [torch.mul, torch.div], ) -def test__mul_tensor_scalar_with_unit_on_gpu( - a: PhlowerDimensionTensor | float, - b: PhlowerDimensionTensor | float, +@given(random_phlower_dimension_tensor(), st.floats()) +def test__op_dimension_tensor_and_float_on_gpu( + op: Callable[ + [PhlowerDimensionTensor, PhlowerDimensionTensor], PhlowerDimensionTensor + ], + a: PhlowerDimensionTensor, + b: float, ): - c = torch.mul(to_cuda_if_needed(a), to_cuda_if_needed(b)) - - assert c == generate_phlower_dimension_tensor().to("cuda:0") - assert c.device == torch.device("cuda:0") + c = op(a, b) - -@pytest.mark.gpu_test -def test__div_tensor_scalar_with_unit_on_gpu(): - a = generate_phlower_dimension_tensor() - b = 2.3 - c = torch.div(to_cuda_if_needed(a), to_cuda_if_needed(b)) - - assert c == generate_phlower_dimension_tensor().to("cuda:0") - assert c.device == torch.device("cuda:0") - - -@pytest.mark.gpu_test -def test__div_scalar_tensor_with_unit_on_gpu(): - a = 2.3 - b = generate_phlower_dimension_tensor() - c = torch.div(to_cuda_if_needed(a), to_cuda_if_needed(b)) - - assert c == generate_phlower_dimension_tensor(inverse=True).to("cuda:0") + cpu_c = op(a.to("cpu"), b) + assert c == cpu_c.to("cuda:0") assert c.device == torch.device("cuda:0") @@ -107,29 +76,14 @@ def test__div_scalar_tensor_with_unit_on_gpu(): "op", [torch.add, torch.sub, torch.mul, torch.div], ) -@pytest.mark.parametrize( - "a, b", - [ - ( - generate_phlower_dimension_tensor(has_dimension=False), - generate_phlower_dimension_tensor(has_dimension=False), - ), - ( - 2.3, - generate_phlower_dimension_tensor(has_dimension=False), - ), - ( - generate_phlower_dimension_tensor(has_dimension=False), - 2.3, - ), +@given(st.floats()) +def test__op_zero_dimension_tensor_on_gpu( + op: Callable[ + [PhlowerDimensionTensor, PhlowerDimensionTensor], PhlowerDimensionTensor ], -) -def test__op_tensor_tensor_without_unit_on_gpu( - op: callable, - a: PhlowerDimensionTensor | float, - b: PhlowerDimensionTensor | float, + x: float, ): - c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) - + a = PhlowerDimensionTensor().to("cuda:0") + c = op(a, x) assert c.device == torch.device("cuda:0") assert c.is_dimensionless diff --git a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py index 2a11e25..e3c725f 100644 --- a/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py +++ b/tests/gpu_tests/test_base/test_tensors/test_phlower_tensor_gpu.py @@ -1,29 +1,64 @@ +from collections.abc import Callable + import numpy as np import pytest import torch +from hypothesis import given +from hypothesis import strategies as st +from hypothesis.extra import numpy as extra_np from phlower import PhlowerTensor, phlower_tensor +from phlower.utils.enums import PhysicalDimensionSymbolType -def generate_random_phlower_tensor( - has_dimension: bool = True, -) -> PhlowerTensor: - if has_dimension: - dimension = {"L": 2, "T": -1} - else: - dimension = {} - return phlower_tensor(torch.rand(5), dimension=dimension) +@st.composite +def random_gpu_phlower_tensors_with_same_dimension( + draw: Callable, shape: tuple[int] | st.SearchStrategy[int] +) -> tuple[PhlowerTensor, PhlowerTensor]: + array1 = draw( + extra_np.arrays( + dtype=np.dtypes.Float32DType(), + shape=shape, + ) + ) + array2 = draw( + extra_np.arrays( + dtype=np.dtypes.Float32DType(), + shape=shape, + ) + ) + + dimensions = draw( + st.lists( + elements=st.floats(allow_nan=False, allow_infinity=False), + min_size=len(PhysicalDimensionSymbolType), + max_size=len(PhysicalDimensionSymbolType), + ) + ) + return ( + phlower_tensor(torch.from_numpy(array1), dimension=dimensions).to( + device="cuda:0" + ), + phlower_tensor(torch.from_numpy(array2), dimension=dimensions).to( + device="cuda:0" + ), + ) -def to_tensor_if_needed(x: PhlowerTensor | float) -> torch.Tensor | float: - if isinstance(x, PhlowerTensor): - return x.to_tensor() - return x +@st.composite +def random_gpu_zerodim_phlower_tensor( + draw: Callable, shape: tuple[int] | st.SearchStrategy[int] +) -> PhlowerTensor: + array = draw( + extra_np.arrays( + dtype=np.dtypes.Float32DType(), + shape=shape, + ) + ) -def to_cuda_if_needed(x: PhlowerTensor | float) -> PhlowerTensor | float: - if isinstance(x, PhlowerTensor): - return x.to("cuda:0") - return x + return phlower_tensor(torch.from_numpy(array), dimension={}).to( + device="cuda:0" + ) @pytest.mark.gpu_test @@ -31,22 +66,16 @@ def to_cuda_if_needed(x: PhlowerTensor | float) -> PhlowerTensor | float: "op", [torch.add, torch.sub, torch.mul, torch.div], ) -@pytest.mark.parametrize( - "a, b", - [ - ( - generate_random_phlower_tensor(), - generate_random_phlower_tensor(), - ), - ], -) -def test__op_tensor_tensor_with_unit_on_gpu( - op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float +@given(random_gpu_phlower_tensors_with_same_dimension(shape=(3, 4))) +def test__op_tensor_with_dimensions_on_gpu( + op: Callable[[PhlowerTensor, PhlowerTensor], PhlowerTensor], + tensors: tuple[PhlowerTensor, PhlowerTensor], ): - c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) - - tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) + a, b = tensors + c = op(a, b) + tc: torch.Tensor = op(a.to_tensor(), b.to_tensor()) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) + assert c.device == torch.device("cuda:0") assert c.dimension.device == torch.device("cuda:0") @@ -54,28 +83,21 @@ def test__op_tensor_tensor_with_unit_on_gpu( @pytest.mark.gpu_test @pytest.mark.parametrize( "op", - [torch.mul, torch.div], + [torch.add, torch.sub, torch.mul, torch.div], ) -@pytest.mark.parametrize( - "a, b", - [ - ( - 2.3, - generate_random_phlower_tensor(), - ), - ( - generate_random_phlower_tensor(), - 2.3, - ), - ], +@given( + random_gpu_zerodim_phlower_tensor(shape=(5, 6)), + st.floats(), ) def test__op_tensor_scalar_with_unit_on_gpu( - op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float + op: Callable[[PhlowerTensor, PhlowerTensor], PhlowerTensor], + a: PhlowerTensor, + b: float, ): - c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) - - tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) + c = op(a, b) + tc: torch.Tensor = op(a.to_tensor(), b) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) + assert c.device == torch.device("cuda:0") assert c.dimension.device == torch.device("cuda:0") @@ -85,29 +107,15 @@ def test__op_tensor_scalar_with_unit_on_gpu( "op", [torch.add, torch.sub, torch.mul, torch.div], ) -@pytest.mark.parametrize( - "a, b", - [ - ( - generate_random_phlower_tensor(has_dimension=False), - generate_random_phlower_tensor(has_dimension=False), - ), - ( - 2.3, - generate_random_phlower_tensor(has_dimension=False), - ), - ( - generate_random_phlower_tensor(has_dimension=False), - 2.3, - ), - ], -) -def test__op_tensor_tensor_without_unit_on_gpu( - op: callable, a: PhlowerTensor | float, b: PhlowerTensor | float +@given(random_gpu_zerodim_phlower_tensor(shape=(10, 5))) +def test__op_tensor_with_zero_dimensions_unit_on_gpu( + op: Callable[[PhlowerTensor, PhlowerTensor], PhlowerTensor], + a: PhlowerTensor, ): - c = op(to_cuda_if_needed(a), to_cuda_if_needed(b)) + b = a.clone() + c = op(a, b) - tc = op(to_tensor_if_needed(a), to_tensor_if_needed(b)) + tc: torch.Tensor = op(a.to_tensor(), b.to_tensor()) np.testing.assert_almost_equal(c.numpy(), tc.cpu().numpy()) assert c.device == torch.device("cuda:0") assert c.dimension.device == torch.device("cuda:0") diff --git a/tests/test_base/test_tensors/test_dimensions.py b/tests/test_base/test_tensors/test_dimensions.py index cc09106..54b87cb 100644 --- a/tests/test_base/test_tensors/test_dimensions.py +++ b/tests/test_base/test_tensors/test_dimensions.py @@ -1,5 +1,7 @@ import pytest import torch +from hypothesis import given +from hypothesis import strategies as st from phlower import PhlowerDimensionTensor from phlower.utils.exceptions import DimensionIncompatibleError @@ -65,3 +67,25 @@ def test__cat_raise_dimension_incompatible(unit1: list[int], unit2: list[int]): with pytest.raises(DimensionIncompatibleError): torch.cat([unit1, unit2], dim=0) + + +@given(st.floats()) +def test__add_with_float_and_non_dimensions(x: float): + dimension = PhlowerDimensionTensor() + + calculated = dimension + x + assert calculated == dimension + + calculated = x + dimension + assert calculated == dimension + + +@given(st.floats()) +def test__mul_with_float_and_non_dimensions(x: float): + dimension = PhlowerDimensionTensor() + + calculated = dimension * x + assert calculated == dimension + + calculated = x * dimension + assert calculated == dimension From 00e7cea02b30552184306e898572b9579da10e60 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Thu, 24 Oct 2024 17:13:12 +0900 Subject: [PATCH 83/89] improve phlower dimensions --- .../_base/tensors/_dimension_tensor.py | 181 +++++++++--------- 1 file changed, 88 insertions(+), 93 deletions(-) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index 4f5cfc0..dc1e8f0 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -5,6 +5,7 @@ from typing import Any import torch +from pipe import uniq from phlower._base._dimension import PhysicalDimensions from phlower.utils.enums import PhysicalDimensionSymbolType @@ -16,16 +17,23 @@ def phlower_dimension_tensor( values: dict[str, float] | PhysicalDimensions, dtype: torch.dtype = torch.float32, + device: str | torch.device = None, ) -> PhlowerDimensionTensor: if not isinstance(values, PhysicalDimensions): values = PhysicalDimensions(values) _list = values.to_list() - return PhlowerDimensionTensor.from_list(_list) + return PhlowerDimensionTensor.from_list(_list, dtype=dtype).to(device) -def zero_dimension_tensor() -> PhlowerDimensionTensor: - return phlower_dimension_tensor({}) +def zero_dimension_tensor( + device: str | torch.device | None, +) -> PhlowerDimensionTensor: + zerodim = phlower_dimension_tensor({}) + if device is None: + return zerodim + + return zerodim.to(device) class PhlowerDimensionTensor: @@ -37,7 +45,9 @@ class PhlowerDimensionTensor: @classmethod def from_list( - cls, values: list[float] | tuple[float] + cls, + values: list[float] | tuple[float], + dtype: torch.dtype = torch.float32, ) -> PhlowerDimensionTensor: """ Parse from list object @@ -55,7 +65,7 @@ def from_list( "length of values is not equal to the number of " f"registered dimension types. input: {len(values)}" ) - _tensor = torch.tensor(values, dtype=torch.float32).reshape( + _tensor = torch.tensor(values, dtype=dtype).reshape( len(PhysicalDimensionSymbolType), 1 ) return PhlowerDimensionTensor(_tensor) @@ -175,69 +185,62 @@ def mean(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: return PhlowerDimensionTensor(inputs._tensor) -def _determine_float_or_dimensions( - inputs: PhlowerDimensionTensor | float, - other: PhlowerDimensionTensor | float, -) -> tuple[float, PhlowerDimensionTensor]: - if isinstance(inputs, float): - assert isinstance( - other, PhlowerDimensionTensor - ), f"one is float, but the other is {other}" - return inputs, other +def _determine_device( + *args: PhlowerDimensionTensor | torch.Tensor | float | int, +) -> torch.device: + devices = {v.device for v in args if isinstance(v, PhlowerDimensionTensor)} - if isinstance(inputs, PhlowerDimensionTensor): - assert isinstance( - other, float - ), f"one is float, but the other is {inputs}" - return other, inputs + if len(devices) != 1: + raise PhlowerDimensionTensor( + f"Cannot determine unique device. {devices}]" + ) - raise ValueError(f"Unexpected situation. inputs: {inputs}, other: {other}") + return devices.pop() + + +def _convert_phlower_dimension_tensors( + *args: PhlowerDimensionTensor | torch.Tensor | float | int, + device: str | torch.device | None = None, +) -> tuple[PhlowerDimensionTensor, ...]: + _dimensions = ( + v + if isinstance(v, PhlowerDimensionTensor) + else zero_dimension_tensor(device=device) + for v in args + ) + return _dimensions @dimension_wrap_implements(torch.add) def add( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor ) -> PhlowerDimensionTensor: - if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): - if inputs != other: - raise DimensionIncompatibleError( - "Add operation for different physical dimensions is not " - "allowed." - ) - - return PhlowerDimensionTensor(inputs._tensor) - - if all( - isinstance(v, (int | float)) or v.is_dimensionless - for v in (inputs, other) - ): - _, dim = _determine_float_or_dimensions(inputs, other) - return zero_dimension_tensor().to(dim.device) + device = _determine_device(inputs) + inputs, other = _convert_phlower_dimension_tensors( + inputs, other, device=device + ) + if inputs != other: + raise DimensionIncompatibleError( + "Add operation for different physical dimensions is not " "allowed." + ) - raise DimensionIncompatibleError() + return PhlowerDimensionTensor(inputs._tensor) @dimension_wrap_implements(torch.sub) def sub( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor ) -> PhlowerDimensionTensor: - if all(isinstance(v, PhlowerDimensionTensor) for v in (inputs, other)): - if inputs != other: - raise DimensionIncompatibleError( - "Sub operation for different physical dimensions is not " - "allowed." - ) - - return PhlowerDimensionTensor(inputs._tensor) - - if all( - isinstance(v, (int | float)) or v.is_dimensionless - for v in (inputs, other) - ): - _, dim = _determine_float_or_dimensions(inputs, other) - return zero_dimension_tensor().to(dim.device) + device = _determine_device(inputs) + inputs, other = _convert_phlower_dimension_tensors( + inputs, other, device=device + ) + if inputs != other: + raise DimensionIncompatibleError( + "Sub operation for different physical dimensions is not " "allowed." + ) - raise DimensionIncompatibleError() + return PhlowerDimensionTensor(inputs._tensor) @dimension_wrap_implements(torch.pow) @@ -257,34 +260,22 @@ def mul( inputs: PhlowerDimensionTensor | torch.Tensor, other: PhlowerDimensionTensor | torch.Tensor, ) -> PhlowerDimensionTensor: - _input = ( - inputs - if isinstance(inputs, PhlowerDimensionTensor) - else zero_dimension_tensor().to(other.device) - ) - _other = ( - other - if isinstance(other, PhlowerDimensionTensor) - else zero_dimension_tensor().to(inputs.device) + device = _determine_device(inputs) + _inputs, _other = _convert_phlower_dimension_tensors( + inputs, other, device=device ) - return PhlowerDimensionTensor(_input._tensor + _other._tensor) + return PhlowerDimensionTensor(_inputs._tensor + _other._tensor) @dimension_wrap_implements(torch.div) def div( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor ) -> PhlowerDimensionTensor: - _input = ( - inputs - if isinstance(inputs, PhlowerDimensionTensor) - else zero_dimension_tensor().to(other.device) - ) - _other = ( - other - if isinstance(other, PhlowerDimensionTensor) - else zero_dimension_tensor().to(inputs.device) + device = _determine_device(inputs) + _inputs, _other = _convert_phlower_dimension_tensors( + inputs, other, device=device ) - return PhlowerDimensionTensor(_input._tensor - _other._tensor) + return PhlowerDimensionTensor(_inputs._tensor - _other._tensor) @dimension_wrap_implements(torch.reshape) @@ -301,12 +292,17 @@ def cat( *, out: PhlowerDimensionTensor = None, ) -> PhlowerDimensionTensor: - if all(isinstance(v, PhlowerDimensionTensor) for v in tensors): - # HACK: is it possible to use unique method ? - for v in tensors: - if v != tensors[0]: - raise DimensionIncompatibleError() - return PhlowerDimensionTensor(tensors[0]._tensor) + device = _determine_device(*tensors) + uniq_inputs: list[PhlowerDimensionTensor] = list( + _convert_phlower_dimension_tensors(*tensors, device=device) | uniq + ) + + if len(uniq_inputs) != 1: + raise DimensionIncompatibleError( + "Only same physical dimensions are allowed when torch.cat." + ) + + return PhlowerDimensionTensor(uniq_inputs[0]._tensor) @dimension_wrap_implements(torch.sparse.mm) @@ -335,15 +331,17 @@ def dropout( @dimension_wrap_implements(torch.stack) def stack(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: - if all(isinstance(v, PhlowerDimensionTensor) for v in inputs): - # HACK: is it possible to use unique method ? - for v in inputs: - if v != inputs[0]: - raise DimensionIncompatibleError() + device = _determine_device(*inputs) + uniq_inputs: list[PhlowerDimensionTensor] = list( + _convert_phlower_dimension_tensors(*inputs, device=device) | uniq + ) - return PhlowerDimensionTensor(inputs[0]._tensor) + if len(uniq_inputs) != 1: + raise DimensionIncompatibleError( + "Only same physical dimensions are allowed when torch.stack" + ) - raise DimensionIncompatibleError() + return PhlowerDimensionTensor(uniq_inputs[0]._tensor) @dimension_wrap_implements(torch.nn.functional.mse_loss) @@ -373,15 +371,12 @@ def _sum( def concatenate( inputs: PhlowerDimensionTensor, *args: Any, **kwards: Any ) -> PhlowerDimensionTensor: - if all(isinstance(v, PhlowerDimensionTensor) for v in inputs): - # HACK: is it possible to use unique method ? - for v in inputs: - if v != inputs[0]: - raise DimensionIncompatibleError() - - return PhlowerDimensionTensor(inputs[0]._tensor) + device = _determine_device(*inputs) + uniq_inputs: list[PhlowerDimensionTensor] = list( + _convert_phlower_dimension_tensors(*inputs, device=device) | uniq + ) - raise DimensionIncompatibleError() + return PhlowerDimensionTensor(uniq_inputs[0]._tensor) @dimension_wrap_implements(torch.tanh) From feba1e5927026233a953ccd74e8d0a5751df3a49 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Thu, 24 Oct 2024 17:48:27 +0900 Subject: [PATCH 84/89] add tests --- .../_base/tensors/_dimension_tensor.py | 30 +++++++++++++++---- src/phlower/_base/tensors/_phlower_tensor.py | 3 +- .../test_base/test_tensors/test_dimensions.py | 27 +++++++++++++++-- 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/phlower/_base/tensors/_dimension_tensor.py b/src/phlower/_base/tensors/_dimension_tensor.py index dc1e8f0..6f06322 100644 --- a/src/phlower/_base/tensors/_dimension_tensor.py +++ b/src/phlower/_base/tensors/_dimension_tensor.py @@ -81,9 +81,17 @@ def __init__( ) return + assert isinstance(tensor, torch.Tensor), f"Unexpected type: {tensor}" + self._tensor = tensor assert self._tensor.shape[0] == len(PhysicalDimensionSymbolType) + def __sub__(self, __value: object): + return torch.sub(self, __value) + + def __rsub__(self, __value: object): + return torch.sub(self, __value) + def __add__(self, __value: object): return torch.add(self, __value) @@ -96,6 +104,12 @@ def __mul__(self, __value: object): def __rmul__(self, __value: object): return torch.mul(self, __value) + def __truediv__(self, __value: object): + return torch.div(self, __value) + + def __rtruediv__(self, __value: object): + return torch.div(__value, self) + def __eq__(self, other: object) -> bool: if not isinstance(other, PhlowerDimensionTensor): return NotImplemented @@ -188,11 +202,15 @@ def mean(inputs: PhlowerDimensionTensor) -> PhlowerDimensionTensor: def _determine_device( *args: PhlowerDimensionTensor | torch.Tensor | float | int, ) -> torch.device: - devices = {v.device for v in args if isinstance(v, PhlowerDimensionTensor)} + devices = { + v.device + for v in args + if isinstance(v, PhlowerDimensionTensor | torch.Tensor) + } if len(devices) != 1: raise PhlowerDimensionTensor( - f"Cannot determine unique device. {devices}]" + f"Cannot determine unique device. {devices}. args: {args}" ) return devices.pop() @@ -215,7 +233,7 @@ def _convert_phlower_dimension_tensors( def add( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor ) -> PhlowerDimensionTensor: - device = _determine_device(inputs) + device = _determine_device(inputs, other) inputs, other = _convert_phlower_dimension_tensors( inputs, other, device=device ) @@ -231,7 +249,7 @@ def add( def sub( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor ) -> PhlowerDimensionTensor: - device = _determine_device(inputs) + device = _determine_device(inputs, other) inputs, other = _convert_phlower_dimension_tensors( inputs, other, device=device ) @@ -260,7 +278,7 @@ def mul( inputs: PhlowerDimensionTensor | torch.Tensor, other: PhlowerDimensionTensor | torch.Tensor, ) -> PhlowerDimensionTensor: - device = _determine_device(inputs) + device = _determine_device(inputs, other) _inputs, _other = _convert_phlower_dimension_tensors( inputs, other, device=device ) @@ -271,7 +289,7 @@ def mul( def div( inputs: PhlowerDimensionTensor, other: PhlowerDimensionTensor ) -> PhlowerDimensionTensor: - device = _determine_device(inputs) + device = _determine_device(inputs, other) _inputs, _other = _convert_phlower_dimension_tensors( inputs, other, device=device ) diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 048963c..7c2dfcb 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -7,6 +7,7 @@ import numpy as np import torch from pipe import select +from typing_extensions import Self from phlower._base._dimension import PhysicalDimensions from phlower._base.tensors._dimension_tensor import ( @@ -248,7 +249,7 @@ def __rtruediv__(self, other: PhlowerTensor) -> PhlowerTensor: def __pow__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.pow(self, other) - def __setitem__(self, key: str, value: float): + def __setitem__(self, key: str, value: float) -> Self: if isinstance(key, PhlowerTensor): self._tensor[key.to_tensor()] = value else: diff --git a/tests/test_base/test_tensors/test_dimensions.py b/tests/test_base/test_tensors/test_dimensions.py index 54b87cb..2a2b070 100644 --- a/tests/test_base/test_tensors/test_dimensions.py +++ b/tests/test_base/test_tensors/test_dimensions.py @@ -69,7 +69,7 @@ def test__cat_raise_dimension_incompatible(unit1: list[int], unit2: list[int]): torch.cat([unit1, unit2], dim=0) -@given(st.floats()) +@given(st.floats(width=32)) def test__add_with_float_and_non_dimensions(x: float): dimension = PhlowerDimensionTensor() @@ -80,7 +80,18 @@ def test__add_with_float_and_non_dimensions(x: float): assert calculated == dimension -@given(st.floats()) +@given(st.floats(width=32)) +def test__sub_with_float_and_non_dimensions(x: float): + dimension = PhlowerDimensionTensor() + + calculated = dimension - x + assert calculated == dimension + + calculated = x - dimension + assert calculated == dimension + + +@given(st.floats(width=32)) def test__mul_with_float_and_non_dimensions(x: float): dimension = PhlowerDimensionTensor() @@ -89,3 +100,15 @@ def test__mul_with_float_and_non_dimensions(x: float): calculated = x * dimension assert calculated == dimension + + +@given(st.floats(width=32)) +def test__div_with_float_and_non_dimensions(x: float): + dimension = PhlowerDimensionTensor() + eps = 1e-5 + + calculated = dimension / (x + eps) + assert calculated == dimension + + calculated = x / dimension + assert calculated == dimension From 424b74700bdfd639f4256e446a9ff741186f6884 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Fri, 25 Oct 2024 23:56:46 +0900 Subject: [PATCH 85/89] add isogcn --- src/phlower/nn/_core_modules/_iso_gcn.py | 335 ++++++++++++++++++ src/phlower/nn/_core_modules/_utils.py | 3 + .../settings/_module_settings/__init__.py | 2 + .../_module_settings/_isogcn_setting.py | 151 ++++++++ 4 files changed, 491 insertions(+) create mode 100644 src/phlower/nn/_core_modules/_iso_gcn.py create mode 100644 src/phlower/settings/_module_settings/_isogcn_setting.py diff --git a/src/phlower/nn/_core_modules/_iso_gcn.py b/src/phlower/nn/_core_modules/_iso_gcn.py new file mode 100644 index 0000000..92b110c --- /dev/null +++ b/src/phlower/nn/_core_modules/_iso_gcn.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +from collections.abc import Callable +from functools import reduce +from typing import Literal + +import torch + +import phlower +from phlower import ISimulationField +from phlower._base.tensors import PhlowerTensor +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn._core_modules import _functions, _utils +from phlower.nn._interface_module import ( + IPhlowerCoreModule, + IReadonlyReferenceGroup, +) +from phlower.settings._module_settings._isogcn_setting import ( + IsoGCNNeumannLinearType, + IsoGCNPropagationType, + IsoGCNSetting, +) + + +class IsoGCN(IPhlowerCoreModule, torch.nn.Module): + """Graph Convolutional Neural Network""" + + @classmethod + def from_setting(cls, setting: IsoGCNSetting) -> IsoGCN: + """Create IsoGCN from IsoGCNSetting instance + + Args: + setting (GCNSetting): setting object for IsoGCN + + Returns: + IsoGCN: IsoGCN object + """ + return IsoGCN(**setting.__dict__) + + @classmethod + def get_nn_name(cls) -> str: + """Return neural network name + + Returns: + str: name + """ + return "IsoGCN" + + @classmethod + def need_reference(cls) -> bool: + return False + + def __init__( + self, + nodes: list[int], + support_names: list[str], + propagations: list[IsoGCNPropagationType], + subchain_activations: list[str], + subchain_dropouts: list[str], + subchain_bias: bool, + coefficient_activations: list[str], + coefficient_dropouts: list[str], + coefficient_bias: bool, + mul_order: Literal["ah_w", "a_hw"] = "ah_w", + neumann_active: bool = False, + neumann_linear_model_type: IsoGCNNeumannLinearType = ( + IsoGCNNeumannLinearType.reuse_graph_weight + ), + neumann_factor: float = 1.0, + neumann_apply_sigmoid_ratio: bool = False, + inversed_moment_name: str = "", + ) -> None: + super().__init__() + + self._weight = _utils.ExtendedLinearList( + nodes=[nodes[0], nodes[-1]], + activations=subchain_activations, + dropouts=subchain_dropouts, + bias=subchain_bias, + ) + self._use_coefficient = True + self._coefficient_network = _utils.ExtendedLinearList( + nodes=nodes, + activations=coefficient_activations, + dropouts=coefficient_dropouts, + bias=coefficient_bias, + ) + self._propagations = propagations + self._nodes = nodes + self._support_names = support_names + self._mul_order = mul_order + + self._inversed_moment_name = inversed_moment_name + self._neumann_active = neumann_active + self._neumann_linear_type = neumann_linear_model_type + self._neumann_factor = neumann_factor + self._neumann_apply_sigmoid_ratio = neumann_apply_sigmoid_ratio + self._neumann_layer = self._create_neumann_layer() + + def _create_neumann_layer(self) -> Callable[[PhlowerTensor], PhlowerTensor]: + if self._neumann_linear_type == IsoGCNNeumannLinearType.identity: + return phlower.nn.functions.identity + + if self._neumann_linear_type == IsoGCNNeumannLinearType.create_new: + return torch.nn.Linear(*self._weight[0].weight.shape, bias=False) + + if ( + self._neumann_linear_type + == IsoGCNNeumannLinearType.reuse_graph_weight + ): + return self._weight[0] + + raise NotImplementedError( + f"{self._neumann_linear_type} is not implemented." + ) + + def resolve( + self, *, parent: IReadonlyReferenceGroup | None = None, **kwards + ) -> None: ... + + def get_reference_name(self) -> str | None: + return None + + def forward( + self, + data: IPhlowerTensorCollections, + *, + field_data: ISimulationField, + **kwards, + ) -> PhlowerTensor: + """forward function which overload torch.nn.Module + + Args: + data (IPhlowerTensorCollections): + data which receives from predecessors + supports (dict[str, PhlowerTensor]): + sparse tensor objects + + Returns: + PhlowerTensor: + Tensor object + """ + assert ( + len(data) <= 2 + ), f"At most two inputs are allowed. input: {len(data)}" + + supports = [field_data[name] for name in self._support_names] + h = self._forward(data[0], supports=supports) + + if self._neumann_active: + h = self._add_neumann( + h, + neumann_condition=data[1], + inversed_moment=field_data[self._inversed_moment_name], + ) + + if not self._use_coefficient: + return h + + h = self._forward_coefficient_network(h) + return h + + def _forward( + self, x: PhlowerTensor, supports: list[PhlowerTensor] + ) -> PhlowerTensor: + if self._mul_order == "ah_w": + h = self._propagate(x, supports) + return self._weight.forward(h) + + if self._mul_order == "a_hw": + h = self._weight.forward(x) + return self._propagate(h) + + raise NotImplementedError( + f"multiplication order: {self._mul_order} is not implemeted." + ) + + def _forward_coefficient_network(self, x: PhlowerTensor) -> PhlowerTensor: + if x.rank() == 0: + coeff = self._coefficient_network.forward(x) + else: + coeff = self._coefficient_network.forward(self._contraction(x)) + + return _functions.einsum( + "i...f,if->i...f", + x, + coeff, + dimension=x.dimension, + is_time_series=x.is_time_series, + is_voxel=x.is_voxel, + ) + + def _propagate( + self, x: PhlowerTensor, supports: list[PhlowerTensor] + ) -> PhlowerTensor: + h = reduce( + lambda y, f: f(y, supports), + [self._select_propagations(name) for name in self._propagations], + initial=x, + ) + return h + + def _add_neumann( + self, + gradient: PhlowerTensor, + neumann_condition: PhlowerTensor, + inversed_moment: PhlowerTensor, + ) -> PhlowerTensor: + neumann_condition = torch.nan_to_num(neumann_condition, nan=0.0) + # NOTE: Shape of inversed_moment is Shape(N, 3, 3, 1) + neumann = ( + _functions.einsum( + "ikl,il...f->ik...f", + inversed_moment[..., 0], + self._neumann_layer(neumann_condition), + dimension=neumann_condition.dimension, + is_time_series=neumann_condition.is_time_series, + is_voxel=neumann_condition.is_voxel, + ) + * self._neumann_factor + ) + if self._neumann_apply_sigmoid_ratio: + sigmoid_coeff = torch.sigmoid(self._coefficient_network[0].weight) + return ( + sigmoid_coeff * gradient + (1.0 - sigmoid_coeff) * neumann + ) * 2 + else: + return gradient + neumann + + def _select_propagations( + self, name: str | IsoGCNPropagationType + ) -> Callable[[PhlowerTensor, list[PhlowerTensor]], PhlowerTensor]: + if name == IsoGCNPropagationType.contraction: + return self._contraction + + if name == IsoGCNPropagationType.convolution: + return self._convolution + + if name == IsoGCNPropagationType.rotation: + return self._rotation + + raise NotImplementedError(f"{name} is not implemented.") + + def _tensor_product( + self, x: PhlowerTensor, supports: list[PhlowerTensor] + ) -> PhlowerTensor: + """Calculate tensor product G \\otimes x. + + Parameters + ---------- + x: PhlowerTensor + [n_vertex, dim, dim, ..., n_feature]-shaped tensor. + supports: list[PhlowerTensor] + List of [n_vertex, n_vertex]-shaped sparse tensor. + + Returns + ------- + y: PhlowerTensor. The rank is incremeted by one from x. + """ + shape = x.shape + if x.rank() < 0: + raise ValueError(f"Tensor shape invalid: {shape}") + + h = torch.stack( + [_functions.spmm(support, x) for support in supports], dim=1 + ) + + return h + + def _convolution( + self, x: PhlowerTensor, supports: list[PhlowerTensor] + ) -> PhlowerTensor: + """Calculate convolution G \\ast x. + + Parameters + ---------- + x: PhlowerTensor + [n_vertex, n_feature]-shaped tensor. + supports: list[torch.Tensor] + List of [n_vertex, n_vertex]-shaped sparse tensor. + + Returns + ------- + y: torch.Tensor + [n_vertex, dim, n_feature]-shaped tensor. + """ + return self._tensor_product(x, supports=supports) + + def _contraction( + self, x: PhlowerTensor, supports: list[PhlowerTensor] + ) -> PhlowerTensor: + """Calculate contraction G \\cdot B. It calculates + \\sum_l G_{i,j,k_1,k_2,...,l} H_{jk_1,k_2,...,l,f} + + Parameters + ---------- + x: PhlowerTensor + [n_vertex, dim, dim, ..., n_feature]-shaped tensor. + supports: list[torch.Tensor] + List of [n_vertex, n_vertex]-shaped sparse tensor. + + Returns + ------- + y: PhlowerTensor + [n_vertex, dim, ..., n_feature]-shaped tensor. + """ + higher_h = self._tensor_product(x, supports=supports) + return _functions.einsum("ikk...->i...", higher_h) + + def _rotation( + self, x: PhlowerTensor, supports: list[PhlowerTensor] + ) -> PhlowerTensor: + """Calculate rotation G \\times x. + + Parameters + ---------- + x: PhlowerTensor + [n_vertex, dim, n_feature]-shaped tensor. + supports: list[PhlowerTensor] + List of [n_vertex, n_vertex]-shaped sparse tensor. + + Returns + ------- + y: PhlowerTensor + [n_vertex, dim, n_feature]-shaped tensor. + """ + h = self._tensor_product(x, supports=supports) + return torch.stack( + [ + h[:, 1, 2] - h[:, 2, 1], + h[:, 2, 0] - h[:, 0, 2], + h[:, 0, 1] - h[:, 1, 0], + ], + dim=1, + ) diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index 78382bd..abb15c9 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -43,6 +43,9 @@ def __init__( def __len__(self) -> int: return len(self._linears) + def __getitem__(self, idx: int) -> torch.nn.Linear: + return self._linears[idx] + def forward_part(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: assert index < self._n_chains diff --git a/src/phlower/settings/_module_settings/__init__.py b/src/phlower/settings/_module_settings/__init__.py index f0fbbb7..995def1 100644 --- a/src/phlower/settings/_module_settings/__init__.py +++ b/src/phlower/settings/_module_settings/__init__.py @@ -10,6 +10,7 @@ ) from phlower.settings._module_settings._gcn_setting import GCNSetting from phlower.settings._module_settings._identity_setting import IdentitySetting +from phlower.settings._module_settings._isogcn_setting import IsoGCNSetting from phlower.settings._module_settings._mlp_setting import MLPSetting from phlower.settings._module_settings._pinv_mlp_setting import PInvMLPSetting from phlower.settings._module_settings._proportional_setting import ( @@ -30,6 +31,7 @@ "Proportional": ProportionalSetting, "Share": ShareSetting, "SimilarityEquivariantMLP": SimilarityEquivariantMLPSetting, + "IsoGCN": IsoGCNSetting, } diff --git a/src/phlower/settings/_module_settings/_isogcn_setting.py b/src/phlower/settings/_module_settings/_isogcn_setting.py new file mode 100644 index 0000000..2d8790d --- /dev/null +++ b/src/phlower/settings/_module_settings/_isogcn_setting.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +from enum import Enum +from typing import Literal + +import pydantic +from pydantic import Field +from typing_extensions import Self + +from phlower.settings._interface import ( + IPhlowerLayerParameters, + IReadOnlyReferenceGroupSetting, +) + + +class IsoGCNPropagationType(Enum, str): + convolution = "convolution" # gradient + contraction = "conttraction" # divergent + rotation = "rotation" + + +class _SubNetworkSetting(pydantic.BaseModel): + is_active: bool = True + activations: list[str] = Field(default_factory=lambda: [], frozen=True) + dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) + bias: bool = Field(False, frozen=True) + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(extra="forbid") + + +class IsoGCNNeumannLinearType(Enum, str): + identity = "identity" + reuse_graph_weight = "reuse_graph_weight" + create_new = "create_new" + + +class _NeumannSetting(pydantic.BaseModel): + is_active: bool = True + linear_model_type: IsoGCNNeumannLinearType = ( + IsoGCNNeumannLinearType.identity + ) + factor: float = 1.0 + apply_sigmoid_ratio: bool = False + inversed_moment_name: str | None = None + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(extra="forbid", frozen=True) + + @pydantic.model_validator(mode="after") + def check_inversed_moment_name_is_not_none(self) -> Self: + if not self.is_active: + return self + + assert ( + self.inversed_moment_name is not None + ), "inversed_moment_name must be determined when using neumann IsoGCN." + + +class IsoGCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): + # This property only overwritten when resolving. + nodes: list[int] = Field(...) + support_names: list[str] = Field(..., frozen=True) + propagations: list[str] = Field(default_factory=lambda: [], frozen=True) + mul_order: Literal["ah_w", "a_hw"] = "ah_w" + graph_weight: _SubNetworkSetting = Field( + default_factory=lambda: _SubNetworkSetting(), frozen=True + ) + coefficient_network: _SubNetworkSetting = Field( + default_factory=lambda: _SubNetworkSetting(), frozen=True + ) + # neumann_options + neumann_setting: _NeumannSetting = Field( + default_factory=lambda: _NeumannSetting(is_active=False), frozen=True + ) + + # special keyward to forbid extra fields in pydantic + model_config = pydantic.ConfigDict(extra="forbid") + + def gather_input_dims(self, *input_dims: int) -> int: + if len(input_dims) != 1: + raise ValueError("only one input is allowed in GCN.") + return input_dims[0] + + @pydantic.field_validator("nodes") + @classmethod + def check_n_nodes(cls, vals: list[int]) -> list[int]: + if len(vals) < 2: + raise ValueError( + "size of nodes must be larger than 1 in IsoGCNSettings." + f" input: {vals}" + ) + + for i, v in enumerate(vals): + if v > 0: + continue + + if (i == 0) and (v == -1): + continue + + raise ValueError( + "nodes in IsoGCN is inconsistent. " + f"value {v} in {i}-th of nodes is not allowed." + ) + + return vals + + @pydantic.model_validator(mode="before") + @classmethod + def fill_empty_activations_dropouts(cls, values: dict) -> dict: + n_nodes = len(values.get("nodes")) + activations = values.get("activations", []) + dropouts = values.get("dropouts", []) + + if len(activations) == 0: + values["activations"] = ["identity" for _ in range(n_nodes - 1)] + + if len(dropouts) == 0: + values["dropouts"] = [0 for _ in range(n_nodes - 1)] + + return values + + @pydantic.model_validator(mode="after") + def check_nodes_size(self) -> Self: + if len(self.nodes) - 1 != len(self.activations): + raise ValueError( + "Size of nodes and activations is not compatible " + "in GCNSettings." + " len(nodes) must be equal to 1 + len(activations)." + ) + + if len(self.nodes) - 1 != len(self.dropouts): + raise ValueError( + "Size of nodes and dropouts is not compatible " + "in GCNSettings." + " len(nodes) must be equal to 1 + len(dropouts)." + ) + return self + + def get_n_nodes(self) -> list[int]: + return self.nodes + + def overwrite_nodes(self, nodes: list[int]) -> None: + self.nodes = nodes + + @property + def need_reference(self) -> bool: + return False + + def get_reference(self, parent: IReadOnlyReferenceGroupSetting) -> None: + return From 8c4c72ece7352b4766d14361345dfbb5f43ed8bf Mon Sep 17 00:00:00 2001 From: sakamoto Date: Mon, 28 Oct 2024 11:11:38 +0900 Subject: [PATCH 86/89] add validation --- src/phlower/nn/_core_modules/_iso_gcn.py | 30 +++++++++++++++++++ src/phlower/nn/_core_modules/_utils.py | 10 +++++++ .../_module_settings/_isogcn_setting.py | 19 ++---------- 3 files changed, 42 insertions(+), 17 deletions(-) diff --git a/src/phlower/nn/_core_modules/_iso_gcn.py b/src/phlower/nn/_core_modules/_iso_gcn.py index 92b110c..16df9b8 100644 --- a/src/phlower/nn/_core_modules/_iso_gcn.py +++ b/src/phlower/nn/_core_modules/_iso_gcn.py @@ -165,9 +165,19 @@ def _forward( ) -> PhlowerTensor: if self._mul_order == "ah_w": h = self._propagate(x, supports) + _validate_rank0_before_applying_nonlinear( + h, + is_bias=self._weight.has_bias(), + activations=self._weight.has_nonlinear_activations(), + ) return self._weight.forward(h) if self._mul_order == "a_hw": + _validate_rank0_before_applying_nonlinear( + x, + is_bias=self._weight.has_bias(), + activations=self._weight.has_nonlinear_activations(), + ) h = self._weight.forward(x) return self._propagate(h) @@ -179,6 +189,7 @@ def _forward_coefficient_network(self, x: PhlowerTensor) -> PhlowerTensor: if x.rank() == 0: coeff = self._coefficient_network.forward(x) else: + # HACK Need to FIX ?? coeff = self._coefficient_network.forward(self._contraction(x)) return _functions.einsum( @@ -333,3 +344,22 @@ def _rotation( ], dim=1, ) + + +def _has_nonlinear_activations(activations: list[str]) -> bool: + return len(v != _functions.identity.__name__ for v in activations) > 0 + + +def _validate_rank0_before_applying_nonlinear( + x: PhlowerTensor, is_bias: bool, activations: list[str] +) -> None: + is_nonlinear_activations = _has_nonlinear_activations(activations) + if x.rank() == 0: + return + + if is_nonlinear_activations or is_bias: + raise ValueError( + "Cannot apply nonlinear operator for rank > 0 tensor." + "Set bias and actications to " + "apply linear operation for rank > 0 tensor" + ) diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index abb15c9..a4687d9 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -25,6 +25,7 @@ def __init__( self._validate_args() self._n_chains = len(self._nodes) + self._is_bias = bias self._linears = torch.nn.ModuleList( [ torch.nn.Linear(n1, n2, bias=bias) @@ -46,6 +47,15 @@ def __len__(self) -> int: def __getitem__(self, idx: int) -> torch.nn.Linear: return self._linears[idx] + def has_bias(self) -> bool: + return self._is_bias + + def has_nonlinear_activations(self) -> bool: + return ( + len(v != _functions.identity.__name__ for v in self._activations) + > 0 + ) + def forward_part(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: assert index < self._n_chains diff --git a/src/phlower/settings/_module_settings/_isogcn_setting.py b/src/phlower/settings/_module_settings/_isogcn_setting.py index 2d8790d..31dca52 100644 --- a/src/phlower/settings/_module_settings/_isogcn_setting.py +++ b/src/phlower/settings/_module_settings/_isogcn_setting.py @@ -13,7 +13,7 @@ ) -class IsoGCNPropagationType(Enum, str): +class IsoGCNPropagationType(str, Enum): convolution = "convolution" # gradient contraction = "conttraction" # divergent rotation = "rotation" @@ -29,7 +29,7 @@ class _SubNetworkSetting(pydantic.BaseModel): model_config = pydantic.ConfigDict(extra="forbid") -class IsoGCNNeumannLinearType(Enum, str): +class IsoGCNNeumannLinearType(str, Enum): identity = "identity" reuse_graph_weight = "reuse_graph_weight" create_new = "create_new" @@ -105,21 +105,6 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: return vals - @pydantic.model_validator(mode="before") - @classmethod - def fill_empty_activations_dropouts(cls, values: dict) -> dict: - n_nodes = len(values.get("nodes")) - activations = values.get("activations", []) - dropouts = values.get("dropouts", []) - - if len(activations) == 0: - values["activations"] = ["identity" for _ in range(n_nodes - 1)] - - if len(dropouts) == 0: - values["dropouts"] = [0 for _ in range(n_nodes - 1)] - - return values - @pydantic.model_validator(mode="after") def check_nodes_size(self) -> Self: if len(self.nodes) - 1 != len(self.activations): From ed39add4724f5988e8a86ffab2ef854d00823202 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Fri, 1 Nov 2024 15:17:24 +0900 Subject: [PATCH 87/89] add tests for IsoGCN. update settings --- src/phlower/_base/tensors/_phlower_tensor.py | 17 +- .../tensors/_tensor_collections.py | 11 +- src/phlower/nn/__init__.py | 1 + src/phlower/nn/_core_modules/_iso_gcn.py | 208 ++++++++------ .../_module_settings/_isogcn_setting.py | 132 ++++++--- .../test_nn/test_core_modules/test_isogcn.py | 269 ++++++++++++++++++ .../isogcn_setting/check_isogcn_nodes.yml | 71 +++++ .../test_isogcn_setting.py | 245 ++++++++++++++++ 8 files changed, 831 insertions(+), 123 deletions(-) create mode 100644 tests/test_nn/test_core_modules/test_isogcn.py create mode 100644 tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_nodes.yml create mode 100644 tests/test_settings/test_module_settings/test_isogcn_setting.py diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 7c2dfcb..6c017cc 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -42,7 +42,7 @@ @overload def phlower_tensor( - tensor: torch.Tensor | PhlowerTensor, + tensor: list | np.ndarray | torch.Tensor | PhlowerTensor, dimension: PhysicDimensionLikeObject | None = None, is_time_series: bool = False, is_voxel: bool = False, @@ -51,14 +51,14 @@ def phlower_tensor( @overload def phlower_tensor( - tensor: torch.Tensor | PhlowerTensor, + tensor: list | np.ndarray | torch.Tensor | PhlowerTensor, dimension: PhysicDimensionLikeObject | None = None, pattern: str = "n...", ) -> PhlowerTensor: ... def phlower_tensor( - tensor: list | torch.Tensor | PhlowerTensor, + tensor: list | np.ndarray | torch.Tensor | PhlowerTensor, dimension: PhysicDimensionLikeObject | None = None, is_time_series: bool | None = None, is_voxel: bool | None = None, @@ -69,7 +69,7 @@ def phlower_tensor( logger.warning("Input dimension_tensor are ignored.") return tensor - if isinstance(tensor, list): + if isinstance(tensor, list | np.ndarray): tensor = torch.tensor(tensor) dimension_tensor = _resolve_dimension_arg(dimension) @@ -282,6 +282,15 @@ def dtype(self) -> torch.dtype: def device(self) -> torch.device: return self._tensor.device + def transpose(self, dim0: int, dim1: int) -> PhlowerTensor: + _tensor = self._tensor.transpose(dim0, dim1) + return PhlowerTensor( + tensor=_tensor, + dimension_tensor=self._dimension_tensor, + is_time_series=self.is_time_series, + is_voxel=self.is_voxel, + ) + def numel(self) -> int: return torch.numel(self._tensor) diff --git a/src/phlower/collections/tensors/_tensor_collections.py b/src/phlower/collections/tensors/_tensor_collections.py index 113e223..88c8883 100644 --- a/src/phlower/collections/tensors/_tensor_collections.py +++ b/src/phlower/collections/tensors/_tensor_collections.py @@ -1,6 +1,7 @@ from __future__ import annotations import abc +from _collections_abc import dict_items from collections.abc import Iterable, Sequence from typing import Any @@ -47,7 +48,7 @@ def keys(self) -> Iterable[str]: ... def values(self): ... @abc.abstractmethod - def pop(self): ... + def pop(self, key: str, default: PhlowerTensor | None = None): ... @abc.abstractmethod def sum(self, weights: dict[str, float] = None) -> PhlowerTensor: ... @@ -95,10 +96,12 @@ def values(self) -> Iterable[dict[str, Any]]: def keys(self) -> Iterable[str]: return self._x.keys() - def items(self) -> abc.ItemsView: + def items(self) -> dict_items[str, PhlowerTensor]: return self._x.items() - def pop(self, key: str, default: PhlowerTensor = None) -> PhlowerTensor: + def pop( + self, key: str, default: PhlowerTensor | None = None + ) -> PhlowerTensor | None: return self._x.pop(key, default) def __getitem__(self, key: str) -> PhlowerTensor: @@ -126,7 +129,7 @@ def min_len(self) -> int: def to_numpy(self) -> dict[str, ArrayDataType]: return { - k: v.tensor().detach().cpu().numpy() for k, v in self._x.items() + k: v.to_tensor().detach().cpu().numpy() for k, v in self.items() } def sum(self, weights: dict[str, float] = None) -> PhlowerTensor: diff --git a/src/phlower/nn/__init__.py b/src/phlower/nn/__init__.py index 8becd39..db01931 100644 --- a/src/phlower/nn/__init__.py +++ b/src/phlower/nn/__init__.py @@ -3,6 +3,7 @@ from phlower.nn._core_modules._en_equivariant_mlp import EnEquivariantMLP from phlower.nn._core_modules._gcn import GCN from phlower.nn._core_modules._identity import Identity +from phlower.nn._core_modules._iso_gcn import IsoGCN from phlower.nn._core_modules._mlp import MLP from phlower.nn._core_modules._pinv_mlp import PInvMLP from phlower.nn._core_modules._proportional import Proportional diff --git a/src/phlower/nn/_core_modules/_iso_gcn.py b/src/phlower/nn/_core_modules/_iso_gcn.py index 16df9b8..5f27359 100644 --- a/src/phlower/nn/_core_modules/_iso_gcn.py +++ b/src/phlower/nn/_core_modules/_iso_gcn.py @@ -1,14 +1,12 @@ from __future__ import annotations from collections.abc import Callable -from functools import reduce from typing import Literal import torch -import phlower -from phlower import ISimulationField from phlower._base.tensors import PhlowerTensor +from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _functions, _utils from phlower.nn._interface_module import ( @@ -16,7 +14,6 @@ IReadonlyReferenceGroup, ) from phlower.settings._module_settings._isogcn_setting import ( - IsoGCNNeumannLinearType, IsoGCNPropagationType, IsoGCNSetting, ) @@ -35,7 +32,25 @@ def from_setting(cls, setting: IsoGCNSetting) -> IsoGCN: Returns: IsoGCN: IsoGCN object """ - return IsoGCN(**setting.__dict__) + return IsoGCN( + nodes=setting.nodes, + isoam_names=setting.isoam_names, + propagations=setting.propagations, + use_self_network=setting.self_network.use_network, + self_network_activations=setting.self_network.activations, + self_network_dropouts=setting.self_network.dropouts, + self_network_bias=setting.self_network.bias, + use_coefficient=setting.coefficient_network.use_network, + coefficient_activations=setting.coefficient_network.activations, + coefficient_dropouts=setting.coefficient_network.dropouts, + coefficient_bias=setting.coefficient_network.bias, + mul_order=setting.mul_order, + to_symmetric=setting.to_symmetric, + neumann_active=setting.neumann_setting.use_neumann, + neumann_factor=setting.neumann_setting.factor, + neumann_input_name=setting.neumann_setting.neumann_input_name, + inversed_moment_name=setting.neumann_setting.inversed_moment_name, + ) @classmethod def get_nn_name(cls) -> str: @@ -52,68 +67,88 @@ def need_reference(cls) -> bool: def __init__( self, - nodes: list[int], - support_names: list[str], + nodes: list[int] | None, + isoam_names: list[str], propagations: list[IsoGCNPropagationType], - subchain_activations: list[str], - subchain_dropouts: list[str], - subchain_bias: bool, - coefficient_activations: list[str], - coefficient_dropouts: list[str], - coefficient_bias: bool, + use_self_network: bool, + self_network_activations: list[str] | None = None, + self_network_dropouts: list[str] | None = None, + self_network_bias: bool | False = False, + use_coefficient: bool = False, + coefficient_activations: list[str] | None = None, + coefficient_dropouts: list[str] | None = None, + coefficient_bias: bool | None = False, mul_order: Literal["ah_w", "a_hw"] = "ah_w", + to_symmetric: bool = False, neumann_active: bool = False, - neumann_linear_model_type: IsoGCNNeumannLinearType = ( - IsoGCNNeumannLinearType.reuse_graph_weight - ), neumann_factor: float = 1.0, - neumann_apply_sigmoid_ratio: bool = False, + neumann_input_name: str | None = None, inversed_moment_name: str = "", ) -> None: super().__init__() - self._weight = _utils.ExtendedLinearList( - nodes=[nodes[0], nodes[-1]], - activations=subchain_activations, - dropouts=subchain_dropouts, - bias=subchain_bias, - ) - self._use_coefficient = True - self._coefficient_network = _utils.ExtendedLinearList( + self._use_self_network = use_self_network + if self._use_self_network: + self._self_network = self._create_layer( + active=use_self_network, + nodes=[nodes[0], nodes[-1]], + activations=self_network_activations, + dropouts=self_network_dropouts, + bias=self_network_bias, + ) + + self._use_coefficient = use_coefficient + self._coefficient_network = self._create_layer( + active=use_coefficient, nodes=nodes, activations=coefficient_activations, dropouts=coefficient_dropouts, bias=coefficient_bias, ) + self._propagations = propagations self._nodes = nodes - self._support_names = support_names + self._isoam_names = isoam_names self._mul_order = mul_order + self._to_symmetric = to_symmetric self._inversed_moment_name = inversed_moment_name self._neumann_active = neumann_active - self._neumann_linear_type = neumann_linear_model_type self._neumann_factor = neumann_factor - self._neumann_apply_sigmoid_ratio = neumann_apply_sigmoid_ratio + self._neumann_input_name = neumann_input_name self._neumann_layer = self._create_neumann_layer() - def _create_neumann_layer(self) -> Callable[[PhlowerTensor], PhlowerTensor]: - if self._neumann_linear_type == IsoGCNNeumannLinearType.identity: - return phlower.nn.functions.identity - - if self._neumann_linear_type == IsoGCNNeumannLinearType.create_new: - return torch.nn.Linear(*self._weight[0].weight.shape, bias=False) - - if ( - self._neumann_linear_type - == IsoGCNNeumannLinearType.reuse_graph_weight - ): - return self._weight[0] - - raise NotImplementedError( - f"{self._neumann_linear_type} is not implemented." + def _create_layer( + self, + active: bool, + nodes: list[int], + activations: list[str], + dropouts: list[float], + bias: bool, + ) -> _utils.ExtendedLinearList | None: + if not active: + return None + + return _utils.ExtendedLinearList( + nodes=nodes, + activations=activations, + dropouts=dropouts, + bias=bias, ) + def _create_neumann_layer( + self, + ) -> Callable[[PhlowerTensor], PhlowerTensor] | None: + if not self._neumann_active: + return None + + if self._self_network is None: + raise ValueError( + "Use self_network when neumannn layer is necessary. " + "It is because neumann layer refers to weight of self_network." + ) + return self._self_network[0] + def resolve( self, *, parent: IReadonlyReferenceGroup | None = None, **kwards ) -> None: ... @@ -144,41 +179,61 @@ def forward( len(data) <= 2 ), f"At most two inputs are allowed. input: {len(data)}" - supports = [field_data[name] for name in self._support_names] - h = self._forward(data[0], supports=supports) + supports = [field_data[name] for name in self._isoam_names] + neumann_value = data.pop(self._neumann_input_name, None) + x = data.unique_item() + + h = self._forward_self_network(x, supports=supports) if self._neumann_active: h = self._add_neumann( h, - neumann_condition=data[1], + neumann_value=neumann_value, inversed_moment=field_data[self._inversed_moment_name], ) + if self._to_symmetric: + # it may be used for velocity + h = (h + h.rearrange("n x1 x2 f -> n x2 x1 f")) / 2.0 + if not self._use_coefficient: return h - h = self._forward_coefficient_network(h) + coeff = self._forward_coefficient_network(x) + _functions.einsum( + "i...f,if->i...f", + h, + coeff, + dimension=h.dimension, + is_time_series=h.is_time_series, + is_voxel=h.is_voxel, + ) return h - def _forward( + def _forward_self_network( self, x: PhlowerTensor, supports: list[PhlowerTensor] ) -> PhlowerTensor: + if not self._use_self_network: + return self._propagate(x, supports) + + assert self._self_network is not None + if self._mul_order == "ah_w": h = self._propagate(x, supports) _validate_rank0_before_applying_nonlinear( h, - is_bias=self._weight.has_bias(), - activations=self._weight.has_nonlinear_activations(), + is_bias=self._self_network.has_bias(), + activations=self._self_network.has_nonlinear_activations(), ) - return self._weight.forward(h) + return self._self_network.forward(h) if self._mul_order == "a_hw": _validate_rank0_before_applying_nonlinear( x, - is_bias=self._weight.has_bias(), - activations=self._weight.has_nonlinear_activations(), + is_bias=self._self_network.has_bias(), + activations=self._self_network.has_nonlinear_activations(), ) - h = self._weight.forward(x) + h = self._self_network.forward(x) return self._propagate(h) raise NotImplementedError( @@ -189,54 +244,38 @@ def _forward_coefficient_network(self, x: PhlowerTensor) -> PhlowerTensor: if x.rank() == 0: coeff = self._coefficient_network.forward(x) else: - # HACK Need to FIX ?? - coeff = self._coefficient_network.forward(self._contraction(x)) + coeff = self._coefficient_network.forward(_functions.contraction(x)) - return _functions.einsum( - "i...f,if->i...f", - x, - coeff, - dimension=x.dimension, - is_time_series=x.is_time_series, - is_voxel=x.is_voxel, - ) + return coeff def _propagate( self, x: PhlowerTensor, supports: list[PhlowerTensor] ) -> PhlowerTensor: - h = reduce( - lambda y, f: f(y, supports), - [self._select_propagations(name) for name in self._propagations], - initial=x, - ) + h = x + for name in self._propagations: + h = self._select_propagations(name)(h, supports) return h def _add_neumann( self, gradient: PhlowerTensor, - neumann_condition: PhlowerTensor, + neumann_value: PhlowerTensor, inversed_moment: PhlowerTensor, ) -> PhlowerTensor: - neumann_condition = torch.nan_to_num(neumann_condition, nan=0.0) + neumann_value = torch.nan_to_num(neumann_value, nan=0.0) # NOTE: Shape of inversed_moment is Shape(N, 3, 3, 1) neumann = ( _functions.einsum( "ikl,il...f->ik...f", inversed_moment[..., 0], - self._neumann_layer(neumann_condition), - dimension=neumann_condition.dimension, - is_time_series=neumann_condition.is_time_series, - is_voxel=neumann_condition.is_voxel, + self._neumann_layer(neumann_value), + dimension=neumann_value.dimension, + is_time_series=neumann_value.is_time_series, + is_voxel=neumann_value.is_voxel, ) * self._neumann_factor ) - if self._neumann_apply_sigmoid_ratio: - sigmoid_coeff = torch.sigmoid(self._coefficient_network[0].weight) - return ( - sigmoid_coeff * gradient + (1.0 - sigmoid_coeff) * neumann - ) * 2 - else: - return gradient + neumann + return gradient + neumann def _select_propagations( self, name: str | IsoGCNPropagationType @@ -250,6 +289,9 @@ def _select_propagations( if name == IsoGCNPropagationType.rotation: return self._rotation + if name == IsoGCNPropagationType.tensor_product: + return self._tensor_product + raise NotImplementedError(f"{name} is not implemented.") def _tensor_product( @@ -363,3 +405,5 @@ def _validate_rank0_before_applying_nonlinear( "Set bias and actications to " "apply linear operation for rank > 0 tensor" ) + + return diff --git a/src/phlower/settings/_module_settings/_isogcn_setting.py b/src/phlower/settings/_module_settings/_isogcn_setting.py index 31dca52..dd59e25 100644 --- a/src/phlower/settings/_module_settings/_isogcn_setting.py +++ b/src/phlower/settings/_module_settings/_isogcn_setting.py @@ -15,12 +15,13 @@ class IsoGCNPropagationType(str, Enum): convolution = "convolution" # gradient - contraction = "conttraction" # divergent + contraction = "contraction" # divergent + tensor_product = "tensor_product" rotation = "rotation" class _SubNetworkSetting(pydantic.BaseModel): - is_active: bool = True + use_network: bool = True activations: list[str] = Field(default_factory=lambda: [], frozen=True) dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) bias: bool = Field(False, frozen=True) @@ -29,62 +30,73 @@ class _SubNetworkSetting(pydantic.BaseModel): model_config = pydantic.ConfigDict(extra="forbid") -class IsoGCNNeumannLinearType(str, Enum): - identity = "identity" - reuse_graph_weight = "reuse_graph_weight" - create_new = "create_new" - - class _NeumannSetting(pydantic.BaseModel): - is_active: bool = True - linear_model_type: IsoGCNNeumannLinearType = ( - IsoGCNNeumannLinearType.identity - ) + use_neumann: bool = True factor: float = 1.0 - apply_sigmoid_ratio: bool = False inversed_moment_name: str | None = None + neumann_input_name: str | None = None # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid", frozen=True) @pydantic.model_validator(mode="after") def check_inversed_moment_name_is_not_none(self) -> Self: - if not self.is_active: + if not self.use_neumann: + return self + + if self.inversed_moment_name is None: + raise ValueError( + "inversed_moment_name must be determined " + "when using neumann IsoGCN." + ) + + return self + + @pydantic.model_validator(mode="after") + def check_valid_neumann_input_name(self) -> Self: + if not self.use_neumann: return self - assert ( - self.inversed_moment_name is not None - ), "inversed_moment_name must be determined when using neumann IsoGCN." + if self.neumann_input_name is None: + raise ValueError("Set neumann input name.") + + return self class IsoGCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): # This property only overwritten when resolving. - nodes: list[int] = Field(...) - support_names: list[str] = Field(..., frozen=True) + nodes: list[int] | None = Field(None) + isoam_names: list[str] = Field(default_factory=lambda: [], frozen=True) propagations: list[str] = Field(default_factory=lambda: [], frozen=True) mul_order: Literal["ah_w", "a_hw"] = "ah_w" - graph_weight: _SubNetworkSetting = Field( - default_factory=lambda: _SubNetworkSetting(), frozen=True + to_symmetric: bool = False + self_network: _SubNetworkSetting = Field( + default_factory=lambda: _SubNetworkSetting(use_network=False), + frozen=True, ) coefficient_network: _SubNetworkSetting = Field( - default_factory=lambda: _SubNetworkSetting(), frozen=True + default_factory=lambda: _SubNetworkSetting(use_network=False), + frozen=True, ) # neumann_options neumann_setting: _NeumannSetting = Field( - default_factory=lambda: _NeumannSetting(is_active=False), frozen=True + default_factory=lambda: _NeumannSetting(use_neumann=False), frozen=True ) # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid") def gather_input_dims(self, *input_dims: int) -> int: - if len(input_dims) != 1: - raise ValueError("only one input is allowed in GCN.") + if (len(input_dims) == 0) or (len(input_dims) >= 3): + raise ValueError("one or two inputs are allowed in IsoGCN.") return input_dims[0] @pydantic.field_validator("nodes") @classmethod - def check_n_nodes(cls, vals: list[int]) -> list[int]: + def check_n_nodes(cls, vals: list[int] | None) -> list[int]: + if vals is None: + return vals + if len(vals) < 2: raise ValueError( "size of nodes must be larger than 1 in IsoGCNSettings." @@ -105,20 +117,74 @@ def check_n_nodes(cls, vals: list[int]) -> list[int]: return vals + @pydantic.field_validator("isoam_names") + @classmethod + def check_isoam_names(cls, vals: list[int]) -> list[int]: + if len(vals) == 0: + raise ValueError("isoam_names is empty.") + + if len(vals) > 3: + raise ValueError("Too many isoam_names is set.") + + return vals + @pydantic.model_validator(mode="after") - def check_nodes_size(self) -> Self: - if len(self.nodes) - 1 != len(self.activations): + def check_valid_self_network(self) -> Self: + if not self.self_network.use_network: + return self + + if len(self.self_network.activations) > 1: raise ValueError( "Size of nodes and activations is not compatible " - "in GCNSettings." - " len(nodes) must be equal to 1 + len(activations)." + "in IsoGCNSettings." + "In self_network, len(activations) <= 1" + f" {self.self_network.activations=}" ) - if len(self.nodes) - 1 != len(self.dropouts): + if len(self.self_network.dropouts) > 1: raise ValueError( "Size of nodes and dropouts is not compatible " - "in GCNSettings." - " len(nodes) must be equal to 1 + len(dropouts)." + "in IsoGCNSettings." + "In self_network, len(activations) <= 1" + f" {self.self_network.dropouts=}" + ) + + return self + + @pydantic.model_validator(mode="after") + def check_valid_coefficient_network(self) -> Self: + if not self.coefficient_network.use_network: + return self + + if len(self.coefficient_network.activations) != 0: + if len(self.coefficient_network.activations) != len(self.nodes) - 1: + raise ValueError( + "Size of nodes and activations is not compatible " + "in IsoGCNSettings." + "In coefficient_network, len(activations) == len(nodes) - 1" + f" {self.coefficient_network.activations=}" + ) + + if len(self.coefficient_network.dropouts) != 0: + if len(self.coefficient_network.dropouts) != len(self.nodes) - 1: + raise ValueError( + "Size of nodes and dropouts is not compatible " + "in IsoGCNSettings." + "In coefficient_network, len(activations) == len(nodes) - 1" + f" {self.coefficient_network.dropouts=}" + ) + + return self + + @pydantic.model_validator(mode="after") + def check_self_weight_when_use_neumann(self) -> Self: + if not self.neumann_setting.use_neumann: + return self + + if not self.self_network.use_network: + raise ValueError( + "Use self_network when neumannn layer is necessary. " + "It is because neumann layer refers to weight of self_network." ) return self diff --git a/tests/test_nn/test_core_modules/test_isogcn.py b/tests/test_nn/test_core_modules/test_isogcn.py new file mode 100644 index 0000000..83a4253 --- /dev/null +++ b/tests/test_nn/test_core_modules/test_isogcn.py @@ -0,0 +1,269 @@ +import numpy as np +import pytest +from phlower import PhlowerTensor, phlower_tensor +from phlower._fields import SimulationField +from phlower.collections import phlower_tensor_collection +from phlower.collections.tensors import IPhlowerTensorCollections +from phlower.nn import IsoGCN + + +def generate_orthogonal_matrix() -> np.ndarray: + def normalize(x: np.ndarray) -> np.ndarray: + return x / np.linalg.norm(x) + + vec1 = normalize(np.random.rand(3) * 2 - 1) + vec2 = normalize(np.random.rand(3) * 2 - 1) + + vec3 = normalize(np.cross(vec1, vec2)) + vec2 = np.cross(vec3, vec1) + return np.array([vec1, vec2, vec3]) + + +def generate_isoam_like_matrix( + location: np.ndarray, key_names: list[str] +) -> IPhlowerTensorCollections: + n_nodes = len(location) + g = np.array( + [ + [location[col] - location[row] for col in range(n_nodes)] + for row in range(n_nodes) + ] + ) + g_eye = np.einsum("ij,ik->ijk", np.eye(n_nodes), np.sum(g, axis=1)) + g_tilde: np.ndarray = g - g_eye + + dict_data: dict[str, PhlowerTensor] = {} + + # print(f"{g_tilde.shape=}") + n_axis = g_tilde.shape[-1] + assert n_axis <= 3 + assert n_axis == len(key_names) + + for i in range(n_axis): + dict_data[key_names[i]] = phlower_tensor(g_tilde[..., i]) + return phlower_tensor_collection(dict_data) + + +class OrthogonalOperator: + def __init__(self): + self._orthogonal = generate_orthogonal_matrix() + + def transform_position(self, location: np.ndarray) -> np.ndarray: + return np.einsum("ij,nj->ni", self._orthogonal, location) + + def transform_rank1(self, array: np.ndarray) -> np.ndarray: + return np.einsum("ip,npk->nik", self._orthogonal, array) + + def transform_rank2(self, array: np.ndarray) -> np.ndarray: + # U X U^T + _tmp = np.einsum("ij,njkl->nikl", self._orthogonal, array) + return np.einsum("njkl,ki->njil", _tmp, self._orthogonal.T) + + +def forward_isogcn_with_no_weight( + location: np.ndarray, propagations: list[str], feature: np.ndarray +) -> PhlowerTensor: + support_names = ["isoam_x", "isoam_y", "isoam_z"] + field = SimulationField( + field_tensors=generate_isoam_like_matrix(location, support_names) + ) + isogcn = IsoGCN( + nodes=[], # This is dummy because no weight is necessary + isoam_names=support_names, + propagations=propagations, + use_self_network=False, + ) + return isogcn.forward( + phlower_tensor_collection({"h": phlower_tensor(feature)}), + field_data=field, + ) + + +# region test for equivariance + + +@pytest.mark.parametrize("n_nodes, n_feature", [(4, 2), (5, 3), (7, 1)]) +def test__equivariance_of_convolution_rank0_to_rank1( + n_nodes: int, n_feature: int +): + # common + propagations = ["convolution"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, propagations=propagations, feature=h + ) + assert h_res.rank() == 1 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = h + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 1 + + # Compare + rotate_h = ortho_op.transform_rank1(h_res.numpy()) + np.testing.assert_array_almost_equal(rotate_h, ortho_h_res.numpy()) + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__equivariance_of_convolution_rank0_to_rank2( + n_nodes: int, n_feature: int +): + # common + propagations = ["convolution", "tensor_product"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, + propagations=propagations, + feature=h, + ) + assert h_res.rank() == 2 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = h + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 2 + + # Compare + # orthogonal transformation + # U X U^T + rotate_h = ortho_op.transform_rank2(h_res.numpy()) + np.testing.assert_array_almost_equal(rotate_h, ortho_h_res.numpy()) + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__equivariance_of_laplacian_rank1_to_rank1( + n_nodes: int, n_feature: int +): + # common + propagations = ["tensor_product", "contraction"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, 3, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, + propagations=propagations, + feature=h, + ) + assert h_res.rank() == 1 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = ortho_op.transform_rank1(h) + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 1 + + # Compare + rotate_h = ortho_op.transform_rank1(h_res.numpy()) + np.testing.assert_array_almost_equal(rotate_h, ortho_h_res.numpy()) + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__invariance_of_contraction_rank1_to_rank0( + n_nodes: int, n_feature: int +): + # common + propagations = ["contraction"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, 3, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, + propagations=propagations, + feature=h, + ) + assert h_res.rank() == 0 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = ortho_op.transform_rank1(h) + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 0 + + # Compare + np.testing.assert_array_almost_equal(h_res.numpy(), ortho_h_res.numpy()) + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__equivariance_of_contraction_rank2_to_rank1( + n_nodes: int, n_feature: int +): + # common + propagations = ["contraction"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, 3, 3, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, + propagations=propagations, + feature=h, + ) + assert h_res.rank() == 1 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = ortho_op.transform_rank2(h) + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 1 + + # Compare + rotate_h = ortho_op.transform_rank1(h_res.numpy()) + np.testing.assert_array_almost_equal(rotate_h, ortho_h_res.numpy()) + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__invariance_of_contraction_rank2_to_rank0( + n_nodes: int, n_feature: int +): + # common + propagations = ["contraction", "contraction"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, 3, 3, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, + propagations=propagations, + feature=h, + ) + assert h_res.rank() == 0 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = ortho_op.transform_rank2(h) + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 0 + + # Compare + np.testing.assert_array_almost_equal(h_res.numpy(), ortho_h_res.numpy()) + + +# endregion diff --git a/tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_nodes.yml b/tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_nodes.yml new file mode 100644 index 0000000..6348dc7 --- /dev/null +++ b/tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_nodes.yml @@ -0,0 +1,71 @@ +misc: + + tests: + IsoGCN0: [12, 20, 5] + GCN1: [5, 20, 5] + GCN2: [5, 5] + +model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: IsoGCN + name: IsoGCN0 + input_keys: + - feature1 + output_key: isogcn0 + destinations: + - GCN1 + nn_parameters: + nodes: [-1, 20, 5] + isoam_names: ["support1"] + self_network: + activations: ["Identity"] + coefficient_network: + activations: ["tanh", "tanh"] + + - nn_type: GCN + name: GCN1 + input_keys: + - isogcn0 + output_key: gcn1 + destinations: + - GCN2 + nn_parameters: + nodes: [-1, 20, 5] + support_name: support1 + + - nn_type: GCN + name: GCN2 + input_keys: + - gcn1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + diff --git a/tests/test_settings/test_module_settings/test_isogcn_setting.py b/tests/test_settings/test_module_settings/test_isogcn_setting.py new file mode 100644 index 0000000..402e704 --- /dev/null +++ b/tests/test_settings/test_module_settings/test_isogcn_setting.py @@ -0,0 +1,245 @@ +import pathlib +from collections.abc import Callable + +import hypothesis.strategies as st +import pytest +import yaml +from hypothesis import assume, given, settings +from phlower.settings import PhlowerModelSetting +from phlower.settings._module_settings import IsoGCNSetting + + +def test__use_flag_when_default_setting(): + _setting = IsoGCNSetting(isoam_names=["dummmy"]) + + assert not _setting.self_network.use_network + assert not _setting.coefficient_network.use_network + assert not _setting.neumann_setting.use_neumann + + +@pytest.mark.parametrize( + "isoam_names", [([]), (["a", "b", "c", "d"]), (["a", "b", "c", "d", "e"])] +) +def test__check_length_of_isoam_names(isoam_names: list[str]): + with pytest.raises(ValueError): + _ = IsoGCNSetting(isoam_names=isoam_names) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["tanh"], [0.2]), + ([10, 30], ["identity"], [0.3]), + ([5, 10, 20, 5], ["relu"], [0.3]), + ([-1, 20, 30], ["tanh"], [0.1]), + ([5, 10, 20, 5], [], []), + ], +) +def test__can_accept_valid_parameters_for_self_network( + nodes: list[int], activations: list[str], dropouts: list[float] +): + _ = IsoGCNSetting( + nodes=nodes, + isoam_names=["dummy"], + self_network={ + "activations": activations, + "dropouts": dropouts, + "bias": True, + }, + ) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["tanh", "tanh"], [0.2]), + ([10, 30], ["identity", "identity"], []), + ([5, 10, 20, 5], ["relu", "relu", "relu"], [0.3, 0.4, 0.0]), + ([-1, 20, 30], [], [0.1, 0.1]), + ], +) +def test__raise_error_invalid_parameters_for_self_network( + nodes: list[int], activations: list[str], dropouts: list[float] +): + with pytest.raises(ValueError): + _ = IsoGCNSetting( + nodes=nodes, + isoam_names=["dummy"], + self_network={ + "activations": activations, + "dropouts": dropouts, + "bias": True, + }, + ) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), + ([10, 30], ["identity"], [0.3]), + ([5, 10, 20, 5], ["relu", "relu", "tanh"], [0.3, 0.2, 0.1]), + ([-1, 20, 30], ["tanh", "tanh"], [0.2, 0.1]), + ], +) +def test__can_accept_valid_parameters_for_coefficient_network( + nodes: list[int], activations: list[str], dropouts: list[float] +): + _ = IsoGCNSetting( + nodes=nodes, + isoam_names=["dummy"], + coefficient_network={ + "activations": activations, + "dropouts": dropouts, + "bias": True, + }, + ) + + +@pytest.mark.parametrize( + "nodes, activations, dropouts", + [ + ([10, 20, 30], ["identity"], []), + ([10, 30], [], [0.3, 0.4]), + ([5, 10, 20, 5], ["relu", "relu", "tanh", "identity"], [0.3, 0.2, 0.1]), + ([5, -1, 20, 5], ["relu", "relu", "tanh"], [0.3, 0.2, 0.1]), + ([5, 10, 20, 5], ["relu", "relu", "tanh"], [0.3]), + ([10], [], []), + ], +) +def test__raise_error_invalid_parameters_for_coefficient_network( + nodes: list[int], activations: list[str], dropouts: list[float] +): + with pytest.raises(ValueError): + _ = IsoGCNSetting( + nodes=nodes, + isoam_names=["dummy"], + coefficient_network={ + "activations": activations, + "dropouts": dropouts, + "bias": True, + }, + ) + + +# region test for Neumann setting + + +def test__raise_error_when_use_neumannn_without_self_network(): + with pytest.raises(ValueError): + _ = IsoGCNSetting( + isoam_names=["dummy"], neumann_setting={"factor": 0.2} + ) + + +@pytest.mark.parametrize( + "inversed_moment_name, neumann_input_name", + [(None, None), ("aaa", None), (None, "aaa")], +) +def test__invalid_neumann_setting( + inversed_moment_name: str, neumann_input_name: str +): + with pytest.raises(ValueError): + _ = IsoGCNSetting( + nodes=[10, 10], + self_network={"use_network": True}, + isoam_names=["dummy"], + neumann_setting={ + "factor": 0.2, + "inversed_moment_name": inversed_moment_name, + "neumann_input_name": neumann_input_name, + }, + ) + + +# endregion + + +# region test for resolve + + +@pytest.mark.parametrize( + "input_dims, desired", + [([30], 30), ([40], 40), ([100], 100), ([20, 20], 20)], +) +def test__gather_input_dims(input_dims: list[int], desired: int): + setting = IsoGCNSetting( + nodes=[10, 20], + isoam_names=["dummy"], + ) + + assert setting.gather_input_dims(*input_dims) == desired + + +@pytest.mark.parametrize("input_dims", [([]), ([40, 400, 10]), ([10, 0, 1])]) +def test__raise_error_invalid_input_dims(input_dims: list[int]): + setting = IsoGCNSetting( + nodes=[10, 20], + isoam_names=["dummy"], + ) + + with pytest.raises(ValueError): + _ = setting.gather_input_dims(*input_dims) + + +@st.composite +def same_length_lists(draw: Callable) -> tuple[list[int], list[int]]: + n_elements = draw(st.integers(min_value=2, max_value=10)) + fixed_length_list = st.lists( + st.integers(min_value=1, max_value=200), + min_size=n_elements, + max_size=n_elements, + ) + + return (draw(fixed_length_list), draw(fixed_length_list)) + + +@given(same_length_lists()) +@settings(max_examples=100) +def test__nodes_is_update_after_overwrite_nodes( + lists: tuple[list[int], list[int]], +): + nodes, update_nodes = lists + assume(nodes != update_nodes) + setting = IsoGCNSetting( + nodes=nodes, + isoam_names=["dummy"], + ) + + before_nodes = setting.get_n_nodes() + assert before_nodes != update_nodes + + setting.overwrite_nodes(update_nodes) + assert setting.get_n_nodes() == update_nodes + + +def test__reference_is_not_necessary(): + setting = IsoGCNSetting(nodes=[10, 20], isoam_names=["dummy"]) + + assert not setting.need_reference + + +# endregion + + +# region E2E tests + +_TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/isogcn_setting" + + +@pytest.mark.parametrize("yaml_file", ["check_isogcn_nodes.yml"]) +def test__nodes_after_resolve(yaml_file: str): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + + assert len(content["misc"]["tests"].items()) > 0 + + for key, value in content["misc"]["tests"].items(): + target = setting.network.search_module_setting(key) + assert target.get_n_nodes() == value + + +# endregion From cbfd605c506677040ca18e0996dc63bbf6509f02 Mon Sep 17 00:00:00 2001 From: sakamoto Date: Sat, 2 Nov 2024 02:27:41 +0900 Subject: [PATCH 88/89] add tests for settings --- src/phlower/_base/tensors/_phlower_tensor.py | 30 +- .../tensors/_unsupported_function_names.py | 4 - src/phlower/nn/_core_modules/_iso_gcn.py | 108 +++--- src/phlower/nn/_core_modules/_utils.py | 63 +++- src/phlower/settings/_group_settings.py | 4 + src/phlower/settings/_interface.py | 15 + .../_module_settings/_concatenator_setting.py | 3 + .../_en_equivariant_mlp_setting.py | 3 + .../settings/_module_settings/_gcn_setting.py | 3 + .../_module_settings/_identity_setting.py | 3 + .../_module_settings/_isogcn_setting.py | 28 ++ .../settings/_module_settings/_mlp_setting.py | 4 + .../_module_settings/_pinv_mlp_setting.py | 4 + .../_module_settings/_proportional_setting.py | 4 + .../_module_settings/_share_setting.py | 4 + .../_similarity_equivariant_mlp_setting.py | 4 + src/phlower/utils/enums.py | 13 + .../test_tensors/test_phlower_tensor.py | 30 +- .../test_nn/test_core_modules/test_isogcn.py | 333 +++++++++++++++++- .../isogcn_setting/check_isogcn_neumann.yml | 80 +++++ .../isogcn_setting/error_missing_neumann.yml | 70 ++++ .../error_not_matched_neumann.yml | 71 ++++ .../error_unnecessary_neumann.yml | 67 ++++ .../test_isogcn_setting.py | 25 +- 24 files changed, 866 insertions(+), 107 deletions(-) delete mode 100644 src/phlower/_base/tensors/_unsupported_function_names.py create mode 100644 tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_neumann.yml create mode 100644 tests/test_settings/test_module_settings/data/isogcn_setting/error_missing_neumann.yml create mode 100644 tests/test_settings/test_module_settings/data/isogcn_setting/error_not_matched_neumann.yml create mode 100644 tests/test_settings/test_module_settings/data/isogcn_setting/error_unnecessary_neumann.yml diff --git a/src/phlower/_base/tensors/_phlower_tensor.py b/src/phlower/_base/tensors/_phlower_tensor.py index 6c017cc..148c7e2 100644 --- a/src/phlower/_base/tensors/_phlower_tensor.py +++ b/src/phlower/_base/tensors/_phlower_tensor.py @@ -1,5 +1,6 @@ from __future__ import annotations +import functools from collections.abc import Callable, Iterable, Sequence from typing import Any, TypeAlias, overload @@ -16,9 +17,6 @@ ) from phlower._base.tensors._interface import IPhlowerTensor from phlower._base.tensors._tensor_shape import PhlowerShapePattern -from phlower._base.tensors._unsupported_function_names import ( - UNSUPPORTED_FUNCTION_NAMES, -) from phlower.utils import get_logger from phlower.utils.exceptions import ( DimensionIncompatibleError, @@ -36,6 +34,11 @@ | tuple[float] ) +_UNSUPPORTED_FUNCTION_NAMES = [ + "einsum", + "reshape", +] + logger = get_logger(__name__) @@ -70,7 +73,7 @@ def phlower_tensor( return tensor if isinstance(tensor, list | np.ndarray): - tensor = torch.tensor(tensor) + tensor = torch.tensor(tensor, dtype=torch.float32) dimension_tensor = _resolve_dimension_arg(dimension) @@ -87,8 +90,8 @@ def phlower_tensor( return PhlowerTensor( tensor=tensor, dimension_tensor=dimension_tensor, - is_time_series=is_time_series, - is_voxel=is_voxel, + is_time_series=bool(is_time_series), + is_voxel=bool(is_voxel), ) @@ -198,7 +201,9 @@ def is_voxel(self) -> bool: def __repr__(self) -> str: return ( f"PhlowerTensor({self._tensor}, " - f"Dimension: {self._dimension_tensor})" + f"Dimension: {self._dimension_tensor}), " + f"is_time_series: {self.is_time_series}, " + f"is_voxel: {self.is_voxel}" ) def __str__(self) -> str: @@ -249,7 +254,14 @@ def __rtruediv__(self, other: PhlowerTensor) -> PhlowerTensor: def __pow__(self, other: PhlowerTensor) -> PhlowerTensor: return torch.pow(self, other) - def __setitem__(self, key: str, value: float) -> Self: + @functools.wraps(torch.Tensor.__getitem__) + def __getitem__(self, key: Any) -> torch.Tensor: + # NOTE: When accessed by index, PhlowerTensor cannot ensure + # its shape pattern. Thus, return only Tensor object + return self._tensor[key] + + @functools.wraps(torch.Tensor.__setitem__) + def __setitem__(self, key: Any, value: Any) -> Self: if isinstance(key, PhlowerTensor): self._tensor[key.to_tensor()] = value else: @@ -428,7 +440,7 @@ def __torch_function__( args: tuple, kwargs: dict | None = None, ) -> PhlowerTensor: - if func.__name__ in UNSUPPORTED_FUNCTION_NAMES: + if func.__name__ in _UNSUPPORTED_FUNCTION_NAMES: raise PhlowerUnsupportedTorchFunctionError( f"Unsupported function: {func.__name__}" ) diff --git a/src/phlower/_base/tensors/_unsupported_function_names.py b/src/phlower/_base/tensors/_unsupported_function_names.py deleted file mode 100644 index c17b45e..0000000 --- a/src/phlower/_base/tensors/_unsupported_function_names.py +++ /dev/null @@ -1,4 +0,0 @@ -UNSUPPORTED_FUNCTION_NAMES = [ - "einsum", - "reshape", -] diff --git a/src/phlower/nn/_core_modules/_iso_gcn.py b/src/phlower/nn/_core_modules/_iso_gcn.py index 5f27359..4f96e16 100644 --- a/src/phlower/nn/_core_modules/_iso_gcn.py +++ b/src/phlower/nn/_core_modules/_iso_gcn.py @@ -5,7 +5,7 @@ import torch -from phlower._base.tensors import PhlowerTensor +from phlower._base.tensors import PhlowerTensor, phlower_tensor from phlower._fields import ISimulationField from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn._core_modules import _functions, _utils @@ -46,7 +46,7 @@ def from_setting(cls, setting: IsoGCNSetting) -> IsoGCN: coefficient_bias=setting.coefficient_network.bias, mul_order=setting.mul_order, to_symmetric=setting.to_symmetric, - neumann_active=setting.neumann_setting.use_neumann, + use_neumann=setting.neumann_setting.use_neumann, neumann_factor=setting.neumann_setting.factor, neumann_input_name=setting.neumann_setting.neumann_input_name, inversed_moment_name=setting.neumann_setting.inversed_moment_name, @@ -80,7 +80,7 @@ def __init__( coefficient_bias: bool | None = False, mul_order: Literal["ah_w", "a_hw"] = "ah_w", to_symmetric: bool = False, - neumann_active: bool = False, + use_neumann: bool = False, neumann_factor: float = 1.0, neumann_input_name: str | None = None, inversed_moment_name: str = "", @@ -89,8 +89,7 @@ def __init__( self._use_self_network = use_self_network if self._use_self_network: - self._self_network = self._create_layer( - active=use_self_network, + self._self_network = _utils.ExtendedLinearList( nodes=[nodes[0], nodes[-1]], activations=self_network_activations, dropouts=self_network_dropouts, @@ -98,13 +97,13 @@ def __init__( ) self._use_coefficient = use_coefficient - self._coefficient_network = self._create_layer( - active=use_coefficient, - nodes=nodes, - activations=coefficient_activations, - dropouts=coefficient_dropouts, - bias=coefficient_bias, - ) + if self._use_coefficient: + self._coefficient_network = _utils.ExtendedLinearList( + nodes=nodes, + activations=coefficient_activations, + dropouts=coefficient_dropouts, + bias=coefficient_bias, + ) self._propagations = propagations self._nodes = nodes @@ -113,33 +112,15 @@ def __init__( self._to_symmetric = to_symmetric self._inversed_moment_name = inversed_moment_name - self._neumann_active = neumann_active + self._use_neumann = use_neumann self._neumann_factor = neumann_factor self._neumann_input_name = neumann_input_name self._neumann_layer = self._create_neumann_layer() - def _create_layer( - self, - active: bool, - nodes: list[int], - activations: list[str], - dropouts: list[float], - bias: bool, - ) -> _utils.ExtendedLinearList | None: - if not active: - return None - - return _utils.ExtendedLinearList( - nodes=nodes, - activations=activations, - dropouts=dropouts, - bias=bias, - ) - def _create_neumann_layer( self, ) -> Callable[[PhlowerTensor], PhlowerTensor] | None: - if not self._neumann_active: + if not self._use_neumann: return None if self._self_network is None: @@ -185,10 +166,10 @@ def forward( h = self._forward_self_network(x, supports=supports) - if self._neumann_active: + if self._use_neumann: h = self._add_neumann( h, - neumann_value=neumann_value, + neumann_tensor=neumann_value, inversed_moment=field_data[self._inversed_moment_name], ) @@ -221,17 +202,13 @@ def _forward_self_network( if self._mul_order == "ah_w": h = self._propagate(x, supports) _validate_rank0_before_applying_nonlinear( - h, - is_bias=self._self_network.has_bias(), - activations=self._self_network.has_nonlinear_activations(), + h, layer=self._self_network ) return self._self_network.forward(h) if self._mul_order == "a_hw": _validate_rank0_before_applying_nonlinear( - x, - is_bias=self._self_network.has_bias(), - activations=self._self_network.has_nonlinear_activations(), + x, layer=self._self_network ) h = self._self_network.forward(x) return self._propagate(h) @@ -259,19 +236,22 @@ def _propagate( def _add_neumann( self, gradient: PhlowerTensor, - neumann_value: PhlowerTensor, + neumann_tensor: PhlowerTensor, inversed_moment: PhlowerTensor, ) -> PhlowerTensor: - neumann_value = torch.nan_to_num(neumann_value, nan=0.0) - # NOTE: Shape of inversed_moment is Shape(N, 3, 3, 1) + neumann_tensor = torch.nan_to_num(neumann_tensor, nan=0.0) + print(f"{neumann_tensor=}") + print(f"{neumann_tensor.is_time_series=}") + print(f"{neumann_tensor.is_voxel=}") + # NOTE: Shape of inversed_moment is Shape(N, 3, 3) neumann = ( _functions.einsum( "ikl,il...f->ik...f", - inversed_moment[..., 0], - self._neumann_layer(neumann_value), - dimension=neumann_value.dimension, - is_time_series=neumann_value.is_time_series, - is_voxel=neumann_value.is_voxel, + inversed_moment, + self._neumann_layer(neumann_tensor), + dimension=neumann_tensor.dimension, + is_time_series=neumann_tensor.is_time_series, + is_voxel=neumann_tensor.is_voxel, ) * self._neumann_factor ) @@ -378,32 +358,34 @@ def _rotation( [n_vertex, dim, n_feature]-shaped tensor. """ h = self._tensor_product(x, supports=supports) - return torch.stack( - [ - h[:, 1, 2] - h[:, 2, 1], - h[:, 2, 0] - h[:, 0, 2], - h[:, 0, 1] - h[:, 1, 0], - ], - dim=1, + return phlower_tensor( + torch.stack( + [ + h[:, 1, 2] - h[:, 2, 1], + h[:, 2, 0] - h[:, 0, 2], + h[:, 0, 1] - h[:, 1, 0], + ], + dim=1, + ), + dimension=h.dimension, + is_time_series=h.is_time_series, + is_voxel=h.is_voxel, ) -def _has_nonlinear_activations(activations: list[str]) -> bool: - return len(v != _functions.identity.__name__ for v in activations) > 0 - - def _validate_rank0_before_applying_nonlinear( - x: PhlowerTensor, is_bias: bool, activations: list[str] + x: PhlowerTensor, layer: _utils.ExtendedLinearList ) -> None: - is_nonlinear_activations = _has_nonlinear_activations(activations) if x.rank() == 0: return - if is_nonlinear_activations or is_bias: + if layer.has_nonlinearity(): raise ValueError( "Cannot apply nonlinear operator for rank > 0 tensor." "Set bias and actications to " - "apply linear operation for rank > 0 tensor" + "apply linear operation for rank > 0 tensor." + f"Layer info: {layer.has_bias()=}, {layer.has_dropout()=}, " + f"{layer.has_nonlinear_activations()=}" ) return diff --git a/src/phlower/nn/_core_modules/_utils.py b/src/phlower/nn/_core_modules/_utils.py index a4687d9..ab11911 100644 --- a/src/phlower/nn/_core_modules/_utils.py +++ b/src/phlower/nn/_core_modules/_utils.py @@ -4,6 +4,7 @@ from phlower._base.tensors import PhlowerTensor from phlower.nn._core_modules import _functions +from phlower.utils.enums import ActivationType from phlower.utils.exceptions import PhlowerInvalidActivationError @@ -47,14 +48,32 @@ def __len__(self) -> int: def __getitem__(self, idx: int) -> torch.nn.Linear: return self._linears[idx] + def has_nonlinearity(self) -> bool: + if self._is_bias: + return True + + if self.has_nonlinear_activations(): + return True + + if self.has_dropout(): + return True + + return False + def has_bias(self) -> bool: return self._is_bias def has_nonlinear_activations(self) -> bool: - return ( - len(v != _functions.identity.__name__ for v in self._activations) - > 0 + n_nonlinear = sum( + 1 for v in self._activations if v != ActivationType.identity ) + return n_nonlinear > 0 + + def has_dropout(self) -> bool: + if len(self._dropouts) == 0: + return False + + return sum(self._dropouts) > 0 def forward_part(self, x: PhlowerTensor, *, index: int) -> PhlowerTensor: assert index < self._n_chains @@ -106,30 +125,38 @@ class ActivationSelector: @staticmethod def select(name: str) -> Callable[[torch.Tensor], torch.Tensor]: - return ActivationSelector._REGISTERED_ACTIVATIONS[name] + type_name = ActivationType[name] + return ActivationSelector._REGISTERED_ACTIVATIONS[type_name.value] @staticmethod def select_inverse( name: str, ) -> Callable[[torch.Tensor], torch.Tensor]: + type_name = ActivationType[name] return ActivationSelector._REGISTERED_ACTIVATIONS[ - ActivationSelector._inverse_activation_name(name) + ActivationSelector._inverse_activation_name(type_name).value ] @staticmethod - def _inverse_activation_name(activation_name: str) -> str: - if activation_name == "identity": - return "identity" - if activation_name == "leaky_relu0p5": - return "inversed_leaky_relu0p5" - if activation_name == "smooth_leaky_relu": - return "inversed_smooth_leaky_relu" - if activation_name == "tanh": - return "truncated_atanh" - - raise PhlowerInvalidActivationError( - f"Cannot inverse for {activation_name}" - ) + def _inverse_activation_name( + activation_name: ActivationType, + ) -> ActivationType: + _to_inverse: dict[ActivationType, ActivationType] = { + ActivationType.identity: ActivationType.identity, + ActivationType.leaky_relu0p5: ActivationType.inversed_leaky_relu0p5, + ActivationType.smooth_leaky_relu: ( + ActivationType.inversed_smooth_leaky_relu + ), + ActivationType.tanh: ActivationType.truncated_atanh, + } + + inverse_type = _to_inverse.get(activation_name) + + if inverse_type is None: + raise PhlowerInvalidActivationError( + f"Cannot inverse for {activation_name}" + ) + return inverse_type @staticmethod def is_exists(name: str) -> bool: diff --git a/src/phlower/settings/_group_settings.py b/src/phlower/settings/_group_settings.py index 93e0b74..a21c6f6 100644 --- a/src/phlower/settings/_group_settings.py +++ b/src/phlower/settings/_group_settings.py @@ -273,6 +273,9 @@ class ModuleSetting(IModuleSetting, pydantic.BaseModel): def get_name(self) -> str: return self.name + def get_input_keys(self) -> list[str]: + return self.input_keys + def get_destinations(self) -> list[str]: return self.destinations @@ -292,6 +295,7 @@ def resolve( _resolved_nodes = self._resolve_nodes(*resolved_outputs) # NOTE: overwrite nodes self.nn_parameters.overwrite_nodes(_resolved_nodes) + self.nn_parameters.confirm(self) def _check_keys(self, *resolved_outputs: dict[str, int]) -> None: _to_input_keys: set[str] = set() diff --git a/src/phlower/settings/_interface.py b/src/phlower/settings/_interface.py index ac11de8..eac7131 100644 --- a/src/phlower/settings/_interface.py +++ b/src/phlower/settings/_interface.py @@ -12,6 +12,9 @@ class IModuleSetting(metaclass=abc.ABCMeta): @abc.abstractmethod def get_name(self) -> str: ... + @abc.abstractmethod + def get_input_keys(self) -> list[str]: ... + @abc.abstractmethod def get_destinations(self) -> list[str]: ... @@ -78,3 +81,15 @@ def get_reference(self, parent: IReadOnlyReferenceGroupSetting): Reference to group setting of its parent """ ... + + @abc.abstractmethod + def confirm(self, self_module: IModuleSetting) -> None: + """Chenck and confirm parameters. This functions is + called after all values of parameters are established. + Write scripts to check its state at last. + + Args: + input_keys (list[str]): keys to input this parameters + **kwards: information of parent module + """ + ... diff --git a/src/phlower/settings/_module_settings/_concatenator_setting.py b/src/phlower/settings/_module_settings/_concatenator_setting.py index ed05309..fe9e704 100644 --- a/src/phlower/settings/_module_settings/_concatenator_setting.py +++ b/src/phlower/settings/_module_settings/_concatenator_setting.py @@ -4,6 +4,7 @@ from pydantic import Field from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -22,6 +23,8 @@ def gather_input_dims(self, *input_dims: int) -> int: sum_dim = sum(v for v in input_dims) return sum_dim + def confirm(self, self_module: IModuleSetting) -> None: ... + @pydantic.field_validator("nodes") @classmethod def check_n_nodes(cls, vals: list[int]) -> list[int]: diff --git a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py index 13069bf..432dc67 100644 --- a/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_en_equivariant_mlp_setting.py @@ -5,6 +5,7 @@ from typing_extensions import Self from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -22,6 +23,8 @@ class EnEquivariantMLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): default_factory=lambda: "identity", frozen=True ) + def confirm(self, self_module: IModuleSetting) -> None: ... + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError("Only one input is allowed in EnEquivariantMLP.") diff --git a/src/phlower/settings/_module_settings/_gcn_setting.py b/src/phlower/settings/_module_settings/_gcn_setting.py index 940c107..015c76a 100644 --- a/src/phlower/settings/_module_settings/_gcn_setting.py +++ b/src/phlower/settings/_module_settings/_gcn_setting.py @@ -5,6 +5,7 @@ from typing_extensions import Self from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -23,6 +24,8 @@ class GCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid") + def confirm(self, self_module: IModuleSetting) -> None: ... + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError("only one input is allowed in GCN.") diff --git a/src/phlower/settings/_module_settings/_identity_setting.py b/src/phlower/settings/_module_settings/_identity_setting.py index 71cf6e7..33fe5cb 100644 --- a/src/phlower/settings/_module_settings/_identity_setting.py +++ b/src/phlower/settings/_module_settings/_identity_setting.py @@ -4,6 +4,7 @@ from pydantic import Field from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -18,6 +19,8 @@ class IdentitySetting(IPhlowerLayerParameters, pydantic.BaseModel): # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid") + def confirm(self, self_module: IModuleSetting) -> None: ... + def gather_input_dims(self, *input_dims: int) -> int: assert len(input_dims) > 0 sum_dim = sum(v for v in input_dims) diff --git a/src/phlower/settings/_module_settings/_isogcn_setting.py b/src/phlower/settings/_module_settings/_isogcn_setting.py index dd59e25..3b0b069 100644 --- a/src/phlower/settings/_module_settings/_isogcn_setting.py +++ b/src/phlower/settings/_module_settings/_isogcn_setting.py @@ -8,6 +8,7 @@ from typing_extensions import Self from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -86,6 +87,33 @@ class IsoGCNSetting(IPhlowerLayerParameters, pydantic.BaseModel): # special keyward to forbid extra fields in pydantic model_config = pydantic.ConfigDict(extra="forbid") + def confirm(self, self_module: IModuleSetting) -> None: + input_keys = self_module.get_input_keys() + if not self.neumann_setting.use_neumann: + if len(input_keys) != 1: + raise ValueError( + "Only one input is allowed " + "when neumann boundary is not considered. " + f"{input_keys=}" + ) + return + + if len(input_keys) != 2: + raise ValueError( + "Two inputs are necessary " + "when neumann boundary is considered. " + f"{input_keys=}" + ) + + if self.neumann_setting.neumann_input_name not in input_keys: + raise ValueError( + f"neumann_input_name is not found in input_keys." + f"neumann_input_name={self.neumann_setting.neumann_input_name}" + f"{input_keys=}" + ) + + return + def gather_input_dims(self, *input_dims: int) -> int: if (len(input_dims) == 0) or (len(input_dims) >= 3): raise ValueError("one or two inputs are allowed in IsoGCN.") diff --git a/src/phlower/settings/_module_settings/_mlp_setting.py b/src/phlower/settings/_module_settings/_mlp_setting.py index 30cdd0b..1a73521 100644 --- a/src/phlower/settings/_module_settings/_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_mlp_setting.py @@ -5,6 +5,7 @@ from typing_extensions import Self from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -17,6 +18,9 @@ class MLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) bias: bool = Field(True, frozen=True) + def confirm(self, self_module: IModuleSetting) -> None: + return + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError("only one input is allowed in MLP.") diff --git a/src/phlower/settings/_module_settings/_pinv_mlp_setting.py b/src/phlower/settings/_module_settings/_pinv_mlp_setting.py index bce9cfc..e5e2cc0 100644 --- a/src/phlower/settings/_module_settings/_pinv_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_pinv_mlp_setting.py @@ -4,6 +4,7 @@ from pydantic import Field from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -19,6 +20,9 @@ class PInvMLPSetting(IPhlowerLayerParameters, pydantic.BaseModel): extra="forbid", arbitrary_types_allowed=True ) + def confirm(self, self_module: IModuleSetting) -> None: + return + def gather_input_dims(self, *input_dims: int) -> int: return self.reference.gather_input_dims(*input_dims) diff --git a/src/phlower/settings/_module_settings/_proportional_setting.py b/src/phlower/settings/_module_settings/_proportional_setting.py index a4ca1fa..27d68a0 100644 --- a/src/phlower/settings/_module_settings/_proportional_setting.py +++ b/src/phlower/settings/_module_settings/_proportional_setting.py @@ -5,6 +5,7 @@ from typing_extensions import Self from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -16,6 +17,9 @@ class ProportionalSetting(IPhlowerLayerParameters, pydantic.BaseModel): ) # This property only overwritten when resolving. dropouts: list[float] = Field(default_factory=lambda: [], frozen=True) + def confirm(self, self_module: IModuleSetting) -> None: + return + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError("Only one input is allowed in EnEquivariantMLP.") diff --git a/src/phlower/settings/_module_settings/_share_setting.py b/src/phlower/settings/_module_settings/_share_setting.py index 808f40d..335d39b 100644 --- a/src/phlower/settings/_module_settings/_share_setting.py +++ b/src/phlower/settings/_module_settings/_share_setting.py @@ -4,6 +4,7 @@ from pydantic import Field from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -19,6 +20,9 @@ class ShareSetting(IPhlowerLayerParameters, pydantic.BaseModel): extra="forbid", arbitrary_types_allowed=True ) + def confirm(self, self_module: IModuleSetting) -> None: + return + def gather_input_dims(self, *input_dims: int) -> int: return self.reference.gather_input_dims(*input_dims) diff --git a/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py b/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py index abc762e..634e0b8 100644 --- a/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py +++ b/src/phlower/settings/_module_settings/_similarity_equivariant_mlp_setting.py @@ -5,6 +5,7 @@ from typing_extensions import Self from phlower.settings._interface import ( + IModuleSetting, IPhlowerLayerParameters, IReadOnlyReferenceGroupSetting, ) @@ -27,6 +28,9 @@ class SimilarityEquivariantMLPSetting( invariant: bool = Field(False, frozen=True) centering: bool = Field(False, frozen=True) + def confirm(self, self_module: IModuleSetting) -> None: + return + def gather_input_dims(self, *input_dims: int) -> int: if len(input_dims) != 1: raise ValueError( diff --git a/src/phlower/utils/enums.py b/src/phlower/utils/enums.py index ade9d15..d14eebb 100644 --- a/src/phlower/utils/enums.py +++ b/src/phlower/utils/enums.py @@ -48,6 +48,19 @@ class PhlowerScalerName(Enum): STD_SCALE = "std_scale" +class ActivationType(str, Enum): + identity = "identity" + inversed_leaky_relu0p5 = "inversed_leaky_relu0p5" + inversed_smooth_leaky_relu = "inversed_smooth_leaky_relu" + leaky_relu0p5 = "leaky_relu0p5" + relu = "relu" + sigmoid = "sigmoid" + smooth_leaky_relu = "smooth_leaky_relu" + sqrt = "sqrt" + tanh = "tanh" + truncated_atanh = "truncated_atanh" + + class PhysicalDimensionSymbolType(Enum): T = 0 # time L = 1 # length diff --git a/tests/test_base/test_tensors/test_phlower_tensor.py b/tests/test_base/test_tensors/test_phlower_tensor.py index 3c8fc7d..f1810ee 100644 --- a/tests/test_base/test_tensors/test_phlower_tensor.py +++ b/tests/test_base/test_tensors/test_phlower_tensor.py @@ -3,6 +3,8 @@ import numpy as np import pytest import torch +from hypothesis import given +from hypothesis import strategies as st from phlower import PhlowerTensor, phlower_dimension_tensor, phlower_tensor from phlower.utils.exceptions import ( DimensionIncompatibleError, @@ -10,6 +12,13 @@ ) +@given(st.lists(st.floats(width=32), min_size=1, max_size=100)) +def test__create_default_phlower_tensor(values: list[float]): + pht = phlower_tensor(values) + assert pht.is_time_series is False + assert pht.is_voxel is False + + def test__create_same_initialized_object_from_list_and_tensor(): list_data = [0.1, 0.2, 0.3] pht_list = phlower_tensor(list_data) @@ -200,11 +209,11 @@ def test__tensor_div_scalar(): c = a / 3.0 ap = PhlowerTensor(a, dims) - cp = ap / 3.0 + cp: PhlowerTensor = ap / 3.0 np.testing.assert_array_almost_equal(cp.to_tensor().numpy(), c.numpy()) - assert cp._dimension_tensor == dims + assert cp.dimension == dims def test__scalar_div_tensor(): @@ -396,3 +405,20 @@ def test__clone(): for k, v in original_dimension_dict.items(): assert cloned.dimension.to_dict()[k] == v assert pht.dimension.to_dict()[k] == 2 * v + + +@pytest.mark.parametrize( + "inputs, nan", + [ + ([0.1, 0.2, float("nan")], 0.0), + ([float("nan"), 0.2, float("nan")], 10.0), + ], +) +def test__nan_to_num(inputs: list[float], nan: float): + tensor = phlower_tensor(inputs) + new_tensor: PhlowerTensor = torch.nan_to_num(tensor, nan=nan) + tensor[torch.isnan(tensor)] = nan + + assert not torch.any(torch.isnan(new_tensor.to_tensor())) + + np.testing.assert_array_almost_equal(new_tensor.numpy(), tensor.numpy()) diff --git a/tests/test_nn/test_core_modules/test_isogcn.py b/tests/test_nn/test_core_modules/test_isogcn.py index 83a4253..6a44748 100644 --- a/tests/test_nn/test_core_modules/test_isogcn.py +++ b/tests/test_nn/test_core_modules/test_isogcn.py @@ -1,10 +1,15 @@ +from typing import Any +from unittest import mock + import numpy as np import pytest from phlower import PhlowerTensor, phlower_tensor from phlower._fields import SimulationField from phlower.collections import phlower_tensor_collection -from phlower.collections.tensors import IPhlowerTensorCollections from phlower.nn import IsoGCN +from phlower.settings._module_settings import IsoGCNSetting + +# region Tools for testing def generate_orthogonal_matrix() -> np.ndarray: @@ -16,18 +21,19 @@ def normalize(x: np.ndarray) -> np.ndarray: vec3 = normalize(np.cross(vec1, vec2)) vec2 = np.cross(vec3, vec1) - return np.array([vec1, vec2, vec3]) + return np.array([vec1, vec2, vec3], dtype=np.float32) def generate_isoam_like_matrix( location: np.ndarray, key_names: list[str] -) -> IPhlowerTensorCollections: +) -> dict[str, PhlowerTensor]: n_nodes = len(location) g = np.array( [ [location[col] - location[row] for col in range(n_nodes)] for row in range(n_nodes) - ] + ], + dtype=np.float32, ) g_eye = np.einsum("ij,ik->ijk", np.eye(n_nodes), np.sum(g, axis=1)) g_tilde: np.ndarray = g - g_eye @@ -41,7 +47,8 @@ def generate_isoam_like_matrix( for i in range(n_axis): dict_data[key_names[i]] = phlower_tensor(g_tilde[..., i]) - return phlower_tensor_collection(dict_data) + + return dict_data class OrthogonalOperator: @@ -61,11 +68,15 @@ def transform_rank2(self, array: np.ndarray) -> np.ndarray: def forward_isogcn_with_no_weight( - location: np.ndarray, propagations: list[str], feature: np.ndarray + location: np.ndarray, + propagations: list[str], + feature: np.ndarray, ) -> PhlowerTensor: support_names = ["isoam_x", "isoam_y", "isoam_z"] field = SimulationField( - field_tensors=generate_isoam_like_matrix(location, support_names) + field_tensors=phlower_tensor_collection( + generate_isoam_like_matrix(location, support_names) + ) ) isogcn = IsoGCN( nodes=[], # This is dummy because no weight is necessary @@ -79,6 +90,117 @@ def forward_isogcn_with_no_weight( ) +# endregion + +# region Test for parsing setting + + +@pytest.mark.parametrize( + "content, desired_items", + [ + ( + { + "isoam_names": ["isoam1"], + "propagations": ["convolution"], + "mul_order": "a_hw", + "to_symmetric": True, + }, + { + "nodes": None, + "isoam_names": ["isoam1"], + "propagations": ["convolution"], + "mul_order": "a_hw", + "to_symmetric": True, + "use_self_network": False, + "use_coefficient": False, + "use_neumann": False, + }, + ), + ( + { + "nodes": [10, 10, 10], + "isoam_names": ["isoam1"], + "propagations": ["contraction"], + "mul_order": "ah_w", + "to_symmetric": False, + "self_network": {"activations": ["identity"], "bias": False}, + "coefficient_network": { + "activations": ["tanh", "tanh"], + "bias": True, + }, + }, + { + "nodes": [10, 10, 10], + "isoam_names": ["isoam1"], + "propagations": ["contraction"], + "mul_order": "ah_w", + "to_symmetric": False, + "use_self_network": True, + "self_network_activations": ["identity"], + "self_network_dropouts": [], + "self_network_bias": False, + "use_coefficient": True, + "coefficient_activations": ["tanh", "tanh"], + "coefficient_dropouts": [], + "coefficient_bias": True, + "use_neumann": False, + }, + ), + ( + { + "nodes": [10, 30], + "isoam_names": ["isoam1"], + "propagations": ["contraction"], + "mul_order": "ah_w", + "to_symmetric": False, + "self_network": { + "activations": ["tanh"], + "dropouts": [0.1], + "bias": False, + }, + "neumann_setting": { + "factor": 0.2, + "neumann_input_name": "neumann_value", + "inversed_moment_name": "inv1", + }, + }, + { + "nodes": [10, 30], + "isoam_names": ["isoam1"], + "propagations": ["contraction"], + "mul_order": "ah_w", + "to_symmetric": False, + "use_self_network": True, + "self_network_activations": ["tanh"], + "self_network_dropouts": [0.1], + "self_network_bias": False, + "use_coefficient": False, + "use_neumann": True, + "neumann_factor": 0.2, + "neumann_input_name": "neumann_value", + "inversed_moment_name": "inv1", + }, + ), + ], +) +def test__passed_correctly_from_setting( + content: dict, desired_items: dict[str, Any] +): + setting = IsoGCNSetting(**content) + + with mock.patch.object(IsoGCN, "__init__") as mocked: + mocked.side_effect = lambda **x: None + _ = IsoGCN.from_setting(setting) + + mocked.assert_called_once() + + kwards: dict = mocked.call_args.kwargs + for k, v in desired_items.items(): + assert kwards.get(k) == v + + +# endregion + # region test for equivariance @@ -141,7 +263,9 @@ def test__equivariance_of_convolution_rank0_to_rank2( # orthogonal transformation # U X U^T rotate_h = ortho_op.transform_rank2(h_res.numpy()) - np.testing.assert_array_almost_equal(rotate_h, ortho_h_res.numpy()) + np.testing.assert_array_almost_equal( + rotate_h, ortho_h_res.numpy(), decimal=5 + ) @pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) @@ -172,7 +296,40 @@ def test__equivariance_of_laplacian_rank1_to_rank1( # Compare rotate_h = ortho_op.transform_rank1(h_res.numpy()) - np.testing.assert_array_almost_equal(rotate_h, ortho_h_res.numpy()) + np.testing.assert_array_almost_equal( + rotate_h, ortho_h_res.numpy(), decimal=5 + ) + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__equivariance_of_rotation_rank1_to_rank1(n_nodes: int, n_feature: int): + # common + propagations = ["rotation"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, 3, n_feature) + h_res = forward_isogcn_with_no_weight( + location=location, + propagations=propagations, + feature=h, + ) + assert h_res.rank() == 1 + + # for location applied by orthogonal matrix + ortho_op = OrthogonalOperator() + ortho_location = ortho_op.transform_position(location) + ortho_h = ortho_op.transform_rank1(h) + ortho_h_res = forward_isogcn_with_no_weight( + location=ortho_location, propagations=propagations, feature=ortho_h + ) + assert ortho_h_res.rank() == 1 + + # Compare + rotate_h = ortho_op.transform_rank1(h_res.numpy()) + np.testing.assert_array_almost_equal( + rotate_h, ortho_h_res.numpy(), decimal=5 + ) @pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) @@ -263,7 +420,163 @@ def test__invariance_of_contraction_rank2_to_rank0( assert ortho_h_res.rank() == 0 # Compare - np.testing.assert_array_almost_equal(h_res.numpy(), ortho_h_res.numpy()) + np.testing.assert_array_almost_equal( + h_res.numpy(), ortho_h_res.numpy(), decimal=5 + ) + + +# endregion + + +# region test for fobidding nonlinear operation +# for tensor of which rank is greater than and equal 1 + + +@pytest.mark.parametrize( + "activations, dropouts, bias", + [(["tanh"], [], False), ([], [], True), ([], [0.1], False)], +) +@pytest.mark.parametrize( + "n_nodes, feature_dims, propagations", + [ + (10, (10, 3, 5), ["convolution"]), + (10, (10, 3, 5), ["contraction"]), + (5, (5, 3, 3, 4), ["contraction"]), + (10, (10, 3, 3, 3, 4), ["contraction", "contraction"]), + ], +) +def test__forbid_operation_for_tensor_with_a_hw( + activations: list[str], + dropouts: list[float], + bias: bool, + n_nodes: int, + feature_dims: tuple[int], + propagations: list[str], +): + location = np.random.rand(n_nodes, 3) + h = np.random.rand(*feature_dims) + + support_names = ["isoam_x", "isoam_y", "isoam_z"] + field = SimulationField( + field_tensors=phlower_tensor_collection( + generate_isoam_like_matrix(location, support_names) + ) + ) + isogcn = IsoGCN( + nodes=[feature_dims[-1], 10], + mul_order="a_hw", + isoam_names=support_names, + propagations=propagations, + use_self_network=True, + self_network_activations=activations, + self_network_dropouts=dropouts, + self_network_bias=bias, + ) + + with pytest.raises(ValueError) as ex: + _ = isogcn.forward( + phlower_tensor_collection({"h": phlower_tensor(h)}), + field_data=field, + ) + + assert "Cannot apply nonlinear operator for rank > 0 tensor." in str(ex) + + +@pytest.mark.parametrize( + "activations, dropouts, bias", + [(["tanh"], [], False), ([], [], True), ([], [0.1], False)], +) +@pytest.mark.parametrize( + "n_nodes, feature_dims, propagations", + [ + (10, (10, 3, 5), ["convolution"]), + (5, (5, 3, 3, 4), ["convolution"]), + (10, (10, 3, 3, 3, 4), ["convolution"]), + ], +) +def test__forbid_operation_for_tensor_with_ah_w( + activations: list[str], + dropouts: list[float], + bias: bool, + n_nodes: int, + feature_dims: tuple[int], + propagations: list[str], +): + location = np.random.rand(n_nodes, 3) + h = np.random.rand(*feature_dims) + + support_names = ["isoam_x", "isoam_y", "isoam_z"] + field = SimulationField( + field_tensors=phlower_tensor_collection( + generate_isoam_like_matrix(location, support_names) + ) + ) + isogcn = IsoGCN( + nodes=[feature_dims[-1], 10], + mul_order="ah_w", + isoam_names=support_names, + propagations=propagations, + use_self_network=True, + self_network_activations=activations, + self_network_dropouts=dropouts, + self_network_bias=bias, + ) + + with pytest.raises(ValueError) as ex: + _ = isogcn.forward( + phlower_tensor_collection({"h": phlower_tensor(h)}), + field_data=field, + ) + + assert "Cannot apply nonlinear operator for rank > 0 tensor." in str(ex) + + +# endregion + +# region Test for Neumann condition + + +@pytest.mark.parametrize("n_nodes, n_feature", [(5, 2), (10, 3)]) +def test__can_forward_with_neumann_condition(n_nodes: int, n_feature: int): + # common + propagations = ["convolution"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, n_feature) + neumann = np.random.rand(n_nodes, 3, n_feature) + + support_names = ["isoam_x", "isoam_y", "isoam_z"] + inversed_moment_name = "inv" + dict_data = generate_isoam_like_matrix(location, support_names) + # NOTE: This inverse moment is dummy data + dict_data.update( + {inversed_moment_name: phlower_tensor(np.random.rand(n_nodes, 3, 3))} + ) + field = SimulationField(field_tensors=phlower_tensor_collection(dict_data)) + + isogcn = IsoGCN( + nodes=[n_feature, 10], + isoam_names=support_names, + propagations=propagations, + use_self_network=True, + self_network_activations=["identity"], + self_network_dropouts=[], + self_network_bias=False, + use_neumann=True, + neumann_input_name="neumann", + inversed_moment_name=inversed_moment_name, + ) + + h_res = isogcn.forward( + phlower_tensor_collection( + {"h": phlower_tensor(h), "neumann": phlower_tensor(neumann)} + ), + field_data=field, + ) + assert h_res.rank() == 1 + +# NOTE: After integrating graphlow, write neumann boudary # endregion diff --git a/tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_neumann.yml b/tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_neumann.yml new file mode 100644 index 0000000..6fd43ac --- /dev/null +++ b/tests/test_settings/test_module_settings/data/isogcn_setting/check_isogcn_neumann.yml @@ -0,0 +1,80 @@ +misc: + + tests: + IsoGCN0: [12, 20, 5] + GCN1: [5, 20, 5] + GCN2: [5, 5] + + +model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + - name: inv + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: neumann_feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: IsoGCN + name: IsoGCN0 + input_keys: + - feature1 + - neumann_feature0 + output_key: isogcn0 + destinations: + - GCN1 + nn_parameters: + nodes: [-1, 20, 5] + isoam_names: ["support1"] + propagations: ["convolution"] + self_network: + activations: ["Identity"] + coefficient_network: + activations: ["tanh", "tanh"] + neumann_setting: + neumann_input_name: neumann_feature0 + inversed_moment_name: "inv" + + - nn_type: GCN + name: GCN1 + input_keys: + - isogcn0 + output_key: gcn1 + destinations: + - GCN2 + nn_parameters: + nodes: [-1, 20, 5] + support_name: support1 + + - nn_type: GCN + name: GCN2 + input_keys: + - gcn1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + diff --git a/tests/test_settings/test_module_settings/data/isogcn_setting/error_missing_neumann.yml b/tests/test_settings/test_module_settings/data/isogcn_setting/error_missing_neumann.yml new file mode 100644 index 0000000..77e7b9b --- /dev/null +++ b/tests/test_settings/test_module_settings/data/isogcn_setting/error_missing_neumann.yml @@ -0,0 +1,70 @@ +model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + - name: inv + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: neumann_feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: IsoGCN + name: IsoGCN0 + input_keys: + - feature1 + output_key: isogcn0 + destinations: + - GCN1 + nn_parameters: + nodes: [-1, 20, 5] + isoam_names: ["support1"] + self_network: + activations: ["Identity"] + coefficient_network: + activations: ["tanh", "tanh"] + neumann_setting: + neumann_input_name: "neumann_feature0" + inversed_moment_name: "inv" + + - nn_type: GCN + name: GCN1 + input_keys: + - isogcn0 + output_key: gcn1 + destinations: + - GCN2 + nn_parameters: + nodes: [-1, 20, 5] + support_name: support1 + + - nn_type: GCN + name: GCN2 + input_keys: + - gcn1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + diff --git a/tests/test_settings/test_module_settings/data/isogcn_setting/error_not_matched_neumann.yml b/tests/test_settings/test_module_settings/data/isogcn_setting/error_not_matched_neumann.yml new file mode 100644 index 0000000..788b1cd --- /dev/null +++ b/tests/test_settings/test_module_settings/data/isogcn_setting/error_not_matched_neumann.yml @@ -0,0 +1,71 @@ +model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + - name: inv + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: neumann_feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: IsoGCN + name: IsoGCN0 + input_keys: + - feature1 + - neumann_feature0 + output_key: isogcn0 + destinations: + - GCN1 + nn_parameters: + nodes: [-1, 20, 5] + isoam_names: ["support1"] + self_network: + activations: ["Identity"] + coefficient_network: + activations: ["tanh", "tanh"] + neumann_setting: + neumann_input_name: "neumann_feature1" + inversed_moment_name: "inv" + + - nn_type: GCN + name: GCN1 + input_keys: + - isogcn0 + output_key: gcn1 + destinations: + - GCN2 + nn_parameters: + nodes: [-1, 20, 5] + support_name: support1 + + - nn_type: GCN + name: GCN2 + input_keys: + - gcn1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + diff --git a/tests/test_settings/test_module_settings/data/isogcn_setting/error_unnecessary_neumann.yml b/tests/test_settings/test_module_settings/data/isogcn_setting/error_unnecessary_neumann.yml new file mode 100644 index 0000000..8dfcf6a --- /dev/null +++ b/tests/test_settings/test_module_settings/data/isogcn_setting/error_unnecessary_neumann.yml @@ -0,0 +1,67 @@ +model: + inputs: + - name: feature0 + members: + - name: feature0 + n_last_dim: 10 + + - name: feature1 + members: + - name: feature1 + n_last_dim: 12 + + fields: + - name: support1 + + network: + nn_type: GROUP + name: CYCLE_MODEL + inputs: + - name: feature0 + n_last_dim: 10 + - name: neumann_feature0 + n_last_dim: 10 + - name: feature1 + n_last_dim: 12 + + outputs: + - name: out_feature0 + n_last_dim: 5 + + modules: + - nn_type: IsoGCN + name: IsoGCN0 + input_keys: + - feature1 + - neumann_feature0 + output_key: isogcn0 + destinations: + - GCN1 + nn_parameters: + nodes: [-1, 20, 5] + isoam_names: ["support1"] + self_network: + activations: ["Identity"] + coefficient_network: + activations: ["tanh", "tanh"] + + - nn_type: GCN + name: GCN1 + input_keys: + - isogcn0 + output_key: gcn1 + destinations: + - GCN2 + nn_parameters: + nodes: [-1, 20, 5] + support_name: support1 + + - nn_type: GCN + name: GCN2 + input_keys: + - gcn1 + output_key: out_feature0 + nn_parameters: + nodes: [-1, 5] + support_name: support1 + diff --git a/tests/test_settings/test_module_settings/test_isogcn_setting.py b/tests/test_settings/test_module_settings/test_isogcn_setting.py index 402e704..c6a8aa2 100644 --- a/tests/test_settings/test_module_settings/test_isogcn_setting.py +++ b/tests/test_settings/test_module_settings/test_isogcn_setting.py @@ -227,7 +227,9 @@ def test__reference_is_not_necessary(): _TEST_DATA_DIR = pathlib.Path(__file__).parent / "data/isogcn_setting" -@pytest.mark.parametrize("yaml_file", ["check_isogcn_nodes.yml"]) +@pytest.mark.parametrize( + "yaml_file", ["check_isogcn_nodes.yml", "check_isogcn_neumann.yml"] +) def test__nodes_after_resolve(yaml_file: str): with open(_TEST_DATA_DIR / yaml_file) as fr: content = yaml.load(fr, Loader=yaml.SafeLoader) @@ -242,4 +244,25 @@ def test__nodes_after_resolve(yaml_file: str): assert target.get_n_nodes() == value +@pytest.mark.parametrize( + "yaml_file, msg", + [ + ("error_missing_neumann.yml", "Two inputs are necessary"), + ("error_unnecessary_neumann.yml", "Only one input is allowed"), + ( + "error_not_matched_neumann.yml", + "neumann_input_name is not found in input_keys", + ), + ], +) +def test__raise_error_for_invalid_setting(yaml_file: str, msg: str): + with open(_TEST_DATA_DIR / yaml_file) as fr: + content = yaml.load(fr, Loader=yaml.SafeLoader) + + with pytest.raises(ValueError) as ex: + setting = PhlowerModelSetting(**content["model"]) + setting.network.resolve(is_first=True) + assert msg in str(ex) + + # endregion From 264fe35d408b23326746ef53895373936128413e Mon Sep 17 00:00:00 2001 From: sakamoto Date: Tue, 5 Nov 2024 10:30:11 +0900 Subject: [PATCH 89/89] add tests for IsoGCN --- src/phlower/nn/_core_modules/_iso_gcn.py | 12 +++-- .../_module_settings/_isogcn_setting.py | 3 -- .../test_nn/test_core_modules/test_isogcn.py | 51 +++++++++++++++++++ 3 files changed, 59 insertions(+), 7 deletions(-) diff --git a/src/phlower/nn/_core_modules/_iso_gcn.py b/src/phlower/nn/_core_modules/_iso_gcn.py index 4f96e16..52f37f0 100644 --- a/src/phlower/nn/_core_modules/_iso_gcn.py +++ b/src/phlower/nn/_core_modules/_iso_gcn.py @@ -95,6 +95,8 @@ def __init__( dropouts=self_network_dropouts, bias=self_network_bias, ) + else: + self._self_network = None self._use_coefficient = use_coefficient if self._use_coefficient: @@ -104,6 +106,8 @@ def __init__( dropouts=coefficient_dropouts, bias=coefficient_bias, ) + else: + self._coefficient_network = None self._propagations = propagations self._nodes = nodes @@ -115,14 +119,14 @@ def __init__( self._use_neumann = use_neumann self._neumann_factor = neumann_factor self._neumann_input_name = neumann_input_name - self._neumann_layer = self._create_neumann_layer() + if self._use_neumann: + self._neumann_layer = self._create_neumann_layer() + else: + self._neumann_layer = None def _create_neumann_layer( self, ) -> Callable[[PhlowerTensor], PhlowerTensor] | None: - if not self._use_neumann: - return None - if self._self_network is None: raise ValueError( "Use self_network when neumannn layer is necessary. " diff --git a/src/phlower/settings/_module_settings/_isogcn_setting.py b/src/phlower/settings/_module_settings/_isogcn_setting.py index 3b0b069..e3a3b44 100644 --- a/src/phlower/settings/_module_settings/_isogcn_setting.py +++ b/src/phlower/settings/_module_settings/_isogcn_setting.py @@ -122,9 +122,6 @@ def gather_input_dims(self, *input_dims: int) -> int: @pydantic.field_validator("nodes") @classmethod def check_n_nodes(cls, vals: list[int] | None) -> list[int]: - if vals is None: - return vals - if len(vals) < 2: raise ValueError( "size of nodes must be larger than 1 in IsoGCNSettings." diff --git a/tests/test_nn/test_core_modules/test_isogcn.py b/tests/test_nn/test_core_modules/test_isogcn.py index 6a44748..4105223 100644 --- a/tests/test_nn/test_core_modules/test_isogcn.py +++ b/tests/test_nn/test_core_modules/test_isogcn.py @@ -577,6 +577,57 @@ def test__can_forward_with_neumann_condition(n_nodes: int, n_feature: int): assert h_res.rank() == 1 +def test__forbid_forward_neumann_wihthout_self_network(): + with pytest.raises(ValueError) as ex: + _ = IsoGCN( + nodes=[10, 10], + isoam_names=["isoam_x"], + propagations=["convolution"], + use_self_network=False, + use_neumann=True, + neumann_input_name="neumann", + inversed_moment_name="inv", + ) + assert "Use self_network when neumannn layer" in str(ex) + + # NOTE: After integrating graphlow, write neumann boudary # endregion + +# region Test for coefficient network + + +@pytest.mark.parametrize("n_nodes, n_feature", [(10, 3), (5, 1), (7, 3)]) +def test__can_forward_with_coefficient_network(n_nodes: int, n_feature: int): + # common + propagations = ["contraction", "contraction"] + + # base coordination + location = np.random.rand(n_nodes, 3) + h = np.random.rand(n_nodes, 3, 3, n_feature) + + support_names = ["isoam_x", "isoam_y", "isoam_z"] + field = SimulationField( + field_tensors=phlower_tensor_collection( + generate_isoam_like_matrix(location, support_names) + ) + ) + isogcn = IsoGCN( + nodes=[n_feature, 5, 5], + isoam_names=support_names, + propagations=propagations, + use_self_network=True, + self_network_activations=["identity"], + use_coefficient=True, + coefficient_activations=["tanh", "tanh"], + ) + h_res = isogcn.forward( + phlower_tensor_collection({"h": phlower_tensor(h)}), + field_data=field, + ) + + assert h_res.rank() == 0 + + +# endregion