Fix missing imports

adendek · adendek · commit 06c244d2d256 · 2026-01-20T16:51:12.000+01:00
diff --git a/pina/__init__.py b/pina/__init__.py
@@ -22,7 +22,3 @@
 from pina._src.core.trainer import Trainer
 from pina._src.condition.condition import Condition
 from pina._src.data.data_module import PinaDataModule
-
-# suppress errors
-import torch._dynamo
-torch._dynamo.config.suppress_errors = True
diff --git a/pina/_src/condition/domain_equation_condition.py b/pina/_src/condition/domain_equation_condition.py
@@ -2,7 +2,7 @@
 
 from pina._src.condition.condition_interface import ConditionInterface
 from pina._src.core.utils import check_consistency
-from pina._src.domain import DomainInterface
+from pina._src.domain.domain_interface import DomainInterface
 from pina._src.equation.equation_interface import EquationInterface
 
 
diff --git a/pina/_src/data/data_module.py b/pina/_src/data/data_module.py
@@ -8,7 +8,7 @@
 from lightning.pytorch import LightningDataModule
 import torch
 from torch_geometric.data import Data
-from torch.utils.data import DataLoader, SequentialSampler, RandomSampler
+from torch.utils.data import DataLoader, SequentialSampler
 from torch.utils.data.distributed import DistributedSampler
 from pina._src.core.label_tensor import LabelTensor
 from pina._src.data.dataset import PinaDatasetFactory, PinaTensorDataset
diff --git a/pina/_src/model/average_neural_operator.py b/pina/_src/model/average_neural_operator.py
@@ -2,9 +2,9 @@
 
 import torch
 from torch import nn
-from .block.average_neural_operator_block import AVNOBlock
-from .kernel_neural_operator import KernelNeuralOperator
-from ..utils import check_consistency
+from pina._src.model.block.average_neural_operator_block import AVNOBlock
+from pina._src.model.kernel_neural_operator import KernelNeuralOperator
+from pina._src.core.utils import check_consistency
 
 
 class AveragingNeuralOperator(KernelNeuralOperator):
diff --git a/pina/_src/model/deeponet.py b/pina/_src/model/deeponet.py
@@ -3,7 +3,7 @@
 from functools import partial
 import torch
 from torch import nn
-from ..utils import check_consistency, is_function
+from pina._src.core.utils import check_consistency, is_function
 
 
 class MIONet(torch.nn.Module):
diff --git a/pina/_src/model/equivariant_graph_neural_operator.py b/pina/_src/model/equivariant_graph_neural_operator.py
@@ -1,8 +1,10 @@
 """Module for the Equivariant Graph Neural Operator model."""
 
 import torch
-from ..utils import check_positive_integer
-from .block.message_passing import EquivariantGraphNeuralOperatorBlock
+from pina._src.core.utils import check_positive_integer
+from pina._src.model.block.message_passing.equivariant_graph_neural_operator_block import (
+    EquivariantGraphNeuralOperatorBlock,
+)
 
 
 class EquivariantGraphNeuralOperator(torch.nn.Module):
diff --git a/pina/_src/model/feed_forward.py b/pina/_src/model/feed_forward.py
@@ -2,8 +2,8 @@
 
 import torch
 from torch import nn
-from ..utils import check_consistency
-from .block.residual import EnhancedLinear
+from pina._src.core.utils import check_consistency
+from pina._src.model.block.residual import EnhancedLinear
 
 
 class FeedForward(torch.nn.Module):
diff --git a/pina/_src/model/fourier_neural_operator.py b/pina/_src/model/fourier_neural_operator.py
@@ -3,10 +3,14 @@
 import warnings
 import torch
 from torch import nn
-from ..label_tensor import LabelTensor
-from ..utils import check_consistency
-from .block.fourier_block import FourierBlock1D, FourierBlock2D, FourierBlock3D
-from .kernel_neural_operator import KernelNeuralOperator
+from pina._src.core.label_tensor import LabelTensor
+from pina._src.core.utils import check_consistency
+from pina._src.model.block.fourier_block import (
+    FourierBlock1D,
+    FourierBlock2D,
+    FourierBlock3D,
+)
+from pina._src.model.kernel_neural_operator import KernelNeuralOperator
 
 
 class FourierIntegralKernel(torch.nn.Module):
diff --git a/pina/_src/model/graph_neural_operator.py b/pina/_src/model/graph_neural_operator.py
@@ -2,8 +2,8 @@
 
 import torch
 from torch.nn import Tanh
-from .block.gno_block import GNOBlock
-from .kernel_neural_operator import KernelNeuralOperator
+from pina._src.model.block.gno_block import GNOBlock
+from pina._src.model.kernel_neural_operator import KernelNeuralOperator
 
 
 class GraphNeuralKernel(torch.nn.Module):
diff --git a/pina/_src/model/kernel_neural_operator.py b/pina/_src/model/kernel_neural_operator.py
@@ -1,7 +1,7 @@
 """Module for the Kernel Neural Operator model class."""
 
 import torch
-from ..utils import check_consistency
+from pina._src.core.utils import check_consistency
 
 
 class KernelNeuralOperator(torch.nn.Module):
diff --git a/pina/_src/model/low_rank_neural_operator.py b/pina/_src/model/low_rank_neural_operator.py
@@ -3,10 +3,10 @@
 import torch
 from torch import nn
 
-from ..utils import check_consistency
+from pina._src.core.utils import check_consistency
 
-from .kernel_neural_operator import KernelNeuralOperator
-from .block.low_rank_block import LowRankBlock
+from pina._src.model.kernel_neural_operator import KernelNeuralOperator
+from pina._src.model.block.low_rank_block import LowRankBlock
 
 
 class LowRankNeuralOperator(KernelNeuralOperator):
diff --git a/pina/_src/model/multi_feed_forward.py b/pina/_src/model/multi_feed_forward.py
@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 import torch
-from .feed_forward import FeedForward
+from pina._src.model.feed_forward import FeedForward
 
 
 class MultiFeedForward(torch.nn.Module, ABC):
diff --git a/pina/_src/model/pirate_network.py b/pina/_src/model/pirate_network.py
@@ -1,8 +1,9 @@
 """Module for the PirateNet model class."""
 
 import torch
-from .block import FourierFeatureEmbedding, PirateNetBlock
-from ..utils import check_consistency, check_positive_integer
+from pina._src.model.block.embedding import FourierFeatureEmbedding
+from pina._src.model.block.pirate_network_block import PirateNetBlock
+from pina._src.core.utils import check_consistency, check_positive_integer
 
 
 class PirateNet(torch.nn.Module):
diff --git a/pina/_src/model/sindy.py b/pina/_src/model/sindy.py
@@ -2,7 +2,7 @@
 
 from typing import Callable
 import torch
-from ..utils import check_consistency, check_positive_integer
+from pina._src.core.utils import check_consistency, check_positive_integer
 
 
 class SINDy(torch.nn.Module):
diff --git a/pina/_src/model/spline.py b/pina/_src/model/spline.py
@@ -2,7 +2,7 @@
 
 import warnings
 import torch
-from ..utils import check_positive_integer, check_consistency
+from pina._src.core.utils import check_consistency, check_positive_integer
 
 
 class Spline(torch.nn.Module):
diff --git a/pina/_src/model/spline_surface.py b/pina/_src/model/spline_surface.py
@@ -2,8 +2,8 @@
 
 import torch
 from .spline import Spline
-from ..label_tensor import LabelTensor
-from ..utils import check_consistency, check_positive_integer
+from pina._src.core.label_tensor import LabelTensor
+from pina._src.core.utils import check_consistency, check_positive_integer
 
 
 class SplineSurface(torch.nn.Module):
diff --git a/pina/_src/problem/abstract_problem.py b/pina/_src/problem/abstract_problem.py
@@ -3,11 +3,14 @@
 from abc import ABCMeta, abstractmethod
 import warnings
 from copy import deepcopy
-from ..utils import check_consistency
-from ..domain import DomainInterface, CartesianDomain
-from ..condition.domain_equation_condition import DomainEquationCondition
-from ..label_tensor import LabelTensor
-from ..utils import merge_tensors, custom_warning_format
+from pina._src.core.utils import check_consistency
+from pina._src.domain.domain_interface import DomainInterface
+from pina._src.domain.cartesian_domain import CartesianDomain
+from pina._src.condition.domain_equation_condition import (
+    DomainEquationCondition,
+)
+from pina._src.core.label_tensor import LabelTensor
+from pina._src.core.utils import merge_tensors, custom_warning_format
 
 
 class AbstractProblem(metaclass=ABCMeta):
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
@@ -5,8 +5,28 @@
 graph-structured data within the PINA training pipeline.
 """
 
-__all__ = ["PinaDataModule", "PinaDataset"]
+from pina._src.data.data_module import (
+    PinaDataModule,
+    PinaSampler,
+    DummyDataloader,
+    Collator,
+    PinaSampler,
+)
 
+from pina._src.data.dataset import (
+    PinaDataset,
+    PinaTensorDataset,
+    PinaGraphDataset,
+    PinaDatasetFactory,
+)
 
-from pina._src.data.data_module import PinaDataModule
-from pina._src.data.dataset import PinaDataset
+__all__ = [
+    "PinaDataModule",
+    "PinaDataset",
+    "PinaSampler",
+    "DummyDataloader",
+    "Collator",
+    "PinaTensorDataset",
+    "PinaGraphDataset",
+    "PinaDatasetFactory",
+]
diff --git a/pina/type_checker.py b/pina/type_checker.py
@@ -0,0 +1,11 @@
+"""Runtime type enforcement and validation utilities.
+
+This module provides decorators to validate function arguments against type
+annotations at runtime. This ensures that PINA components receive inputs
+adhering to expected specifications, improving the robustness of the
+scientific computing pipeline.
+"""
+
+from pina._src.core.type_checker import enforce_types
+
+__all__ = ["enforce_types"]
diff --git a/pina/utils.py b/pina/utils.py
@@ -0,0 +1,19 @@
+"""
+Utility functions for tensor manipulation and input validation.
+
+This module provides helper functions to manage tensor operations and ensure
+data consistency across the PINA framework, supporting robust input checking
+and seamless data merging.
+"""
+
+from pina._src.core.utils import (
+    merge_tensors,
+    check_consistency,
+    check_positive_integer,
+)
+
+__all__ = [
+    "merge_tensors",
+    "check_consistency",
+    "check_positive_integer",
+]
diff --git a/tests/test_block/test_pod.py b/tests/test_block/test_pod.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 
-from pina.model.block.pod_block import PODBlock
+from pina.model.block import PODBlock
 
 x = torch.linspace(-1, 1, 100)
 toy_snapshots = torch.vstack(
diff --git a/tests/test_block/test_rbf.py b/tests/test_block/test_rbf.py
@@ -2,7 +2,7 @@
 import pytest
 import math
 
-from pina.model.block.rbf_block import RBFBlock
+from pina.model.block import RBFBlock
 
 x = torch.linspace(-1, 1, 100)
 toy_params = torch.linspace(0, 1, 10).unsqueeze(1)
diff --git a/tests/test_condition.py b/tests/test_condition.py
@@ -18,7 +18,7 @@
     GraphDataCondition,
 )
 from pina.domain import CartesianDomain
-from pina.equation.equation_factory import FixedValue
+from pina.equation import FixedValue
 from pina.graph import RadiusGraph
 
 example_domain = CartesianDomain({"x": [0, 1], "y": [0, 1]})
diff --git a/tests/test_data/test_data_module.py b/tests/test_data/test_data_module.py
@@ -1,10 +1,10 @@
 import torch
 import pytest
 from pina.data import PinaDataModule
-from pina.data.dataset import PinaTensorDataset, PinaGraphDataset
+from pina.data import PinaTensorDataset, PinaGraphDataset
 from pina.problem.zoo import SupervisedProblem
 from pina.graph import RadiusGraph
-from pina.data.data_module import DummyDataloader
+from pina.data import DummyDataloader
 from pina import Trainer
 from pina.solver import SupervisedSolver
 from torch_geometric.data import Batch
diff --git a/tests/test_data/test_graph_dataset.py b/tests/test_data/test_graph_dataset.py
@@ -1,6 +1,6 @@
 import torch
 import pytest
-from pina.data.dataset import PinaDatasetFactory, PinaGraphDataset
+from pina.data import PinaDatasetFactory, PinaGraphDataset
 from pina.graph import KNNGraph
 from torch_geometric.data import Data
 
diff --git a/tests/test_data/test_tensor_dataset.py b/tests/test_data/test_tensor_dataset.py
@@ -1,6 +1,6 @@
 import torch
 import pytest
-from pina.data.dataset import PinaDatasetFactory, PinaTensorDataset
+from pina.data import PinaDatasetFactory, PinaTensorDataset
 
 input_tensor = torch.rand((100, 10))
 output_tensor = torch.rand((100, 2))
diff --git a/tests/test_loss/test_power_loss.py b/tests/test_loss/test_power_loss.py
@@ -1,5 +1,4 @@
 import torch
-import pytest
 
 from pina.loss import PowerLoss
 
diff --git a/tests/test_problem_zoo/test_supervised_problem.py b/tests/test_problem_zoo/test_supervised_problem.py
@@ -1,7 +1,7 @@
 import torch
 from pina.problem import AbstractProblem
 from pina.condition import InputTargetCondition
-from pina.problem.zoo.supervised_problem import SupervisedProblem
+from pina.problem.zoo import SupervisedProblem
 from pina.graph import RadiusGraph
 
 
