Make ResSim dataclass.

Use __setattr__ instead of @Property

TPFA_ResSim/__init__.py

@@ -1,5 +1,6 @@
 """.. include:: README.md"""
 
+from dataclasses import dataclass
 import numpy as np
 from scipy import sparse
 from scipy.sparse.linalg import spsolve
@@ -10,6 +11,7 @@
 from TPFA_ResSim.plotting import Plot2D
 
 
+@dataclass
 class ResSim(NicePrint, Grid2D, Plot2D):
  """Reservoir simulator class.
 
@@ -32,88 +34,61 @@ class ResSim(NicePrint, Grid2D, Plot2D):
  >>> S[-1, [100, 1300, 2900]]
  array([0.9429345 , 0.91358172, 0.71554613])
  """
- def __init__(self, **kwargs):
- """Write keyword arguments to `self`."""
-
- # Init super
- super().__init__(**{k: kwargs.pop(k) for k in list(kwargs)
- if k in ["Lx", "Ly", "Nx", "Ny"]})
-
- # Add defaults
- defaults = dict(K=np.ones((2, *self.shape)), por=np.ones(self.shape),
- vw=1., vo=1., swc=0., sor=0.)
- kwargs = dict(defaults, **kwargs)
-
- # Write
- for k in list(kwargs):
- setattr(self, k, kwargs.pop(k))
-
-
- K: np.ndarray
+ # Dont use dataclass repr
+ __repr__ = NicePrint.__repr__
+ __str__ = NicePrint.__str__
+
+ def __post_init__(self):
+ defaults = dict(K=np.ones((2, *self.shape)),
+ por=np.ones(self.shape))
+ for k, v in defaults.items():
+ if getattr(self, k) is None:
+ setattr(self, k, v)
+
+ def __setattr__(self, key, val):
+ # Well positions -- collocate at some node
+ if key in ["inj_xy", "prod_xy"] and val is not None:
+ val = np.array(val, float)
+ for i, (x, y) in enumerate(val):
+ val[i] = self.ind2xy(self.xy2ind(x, y))
+ # Well rates
+ if key in ["inj_rates", "prod_rates"] and val is not None:
+ val = np.array(val, float)
+ # Set
+ super().__setattr__(key, val)
+
+ vw: float = 1.
+ """Viscosity for water."""
+ vo: float = 1.
+ """Viscosity for oil."""
+ swc: float = 0.
+ """Irreducible saturation, water."""
+ sor: float = 0.
+ """Irreducible saturation, oil."""
+ K: np.ndarray = None
  """Permeabilities (in x and y directions). Array of shape `(2, Nx, Ny)`)."""
- por: np.ndarray
+ por: np.ndarray = None
  """Porosity; Array of shape `(Nx, Ny)`)."""
- vw: float
- """Viscosity for water. Defaults: 1."""
- vo: float
- """Viscosity for oil. Defaults: 1."""
- swc: float
- """Irreducible saturation, water. Default: 0"""
- sor: float
- """Irreducible saturation, oil. Default: 0"""
-
- @property
- def inj_xy(self):
- """Array of shape `(n, 2)` of x- and y-coords for `n` injector wells.
-
- Values should be betwen `0` and `Lx` or `Ly`.
-
- .. warning:: The wells get co-located with grid nodes, not distributed over nearby ones.
- This is a design choice, not a mathematical necessity.
- """
- return self._inj_xy
-
- @property
- def prod_xy(self):
- """Like `inj_xy`, but for producing wells."""
- return self._prod_xy
-
- def _collocate_at_node(self, xys):
- xys = np.array(xys, float)
- for i, (x, y) in enumerate(xys):
- xys[i] = self.ind2xy(self.xy2ind(x, y))
- return xys
-
- @inj_xy.setter
- def inj_xy(self, value):
- self._inj_xy = self._collocate_at_node(value)
-
- @prod_xy.setter
- def prod_xy(self, value):
- self._prod_xy = self._collocate_at_node(value)
-
- @property
- def inj_rates(self):
- """Array of shape `(n, nTime)` -- or `(n, 1)` if constant-in-time.
-
- .. note:: Both `inj_rates` and `prod_rates` are rates should be positive.
- At each time index, it is asserted that the rates sum to 0,
- otherwise the model would silently input deficit from SW corner.
- """
- return self._inj_rates
- @property
- def prod_rates(self):
- """Like `prod_rates`, but for producing wells."""
- return self._prod_rates
 
- @inj_rates.setter
- def inj_rates(self, value):
- self._inj_rates = np.array(value, float)
+ inj_xy: np.ndarray = None
+ """Array of shape `(n, 2)` of x- and y-coords for `n` injector wells.
 
- @prod_rates.setter
- def prod_rates(self, value):
- self._prod_rates = np.array(value, float)
+ Values should be betwen `0` and `Lx` or `Ly`.
 
+ .. warning:: The wells get co-located with grid nodes, not distributed over nearby ones.
+ This is a design choice, not a mathematical necessity.
+ """
+ prod_xy: np.ndarray = None
+ """Like `inj_xy`, but for producing wells."""
+ inj_rates: np.ndarray = None
+ """Array of shape `(n, nTime)` -- or `(n, 1)` if constant-in-time.
+
+ .. note:: Both `inj_rates` and `prod_rates` are rates should be positive.
+ At each time index, it is asserted that the rates sum to 0,
+ otherwise the model would silently input deficit from SW corner.
+ """
+ prod_rates: np.ndarray = None
+ """Like `prod_rates`, but for producing wells."""
 
  def _set_Q(self, k):
  """Define source/sink *field* (at time `k`) from well specifications."""

TPFA_ResSim/grid.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 
-@dataclass(frozen=True)
+@dataclass
 class Grid2D:
  """Defines a 2D rectangular grid.