subclassing ndarray

lyncs_quda/clover_field.py

@@ -28,95 +28,96 @@ class CloverField(LatticeField):
      This is designed as an intermediary to QUDA CloverField class
      so that it should have 1-to-1 correspondence to an QUDA instance.
     Note:
-     * This class stores the corresponding gauge field in its "field" attribute
-        to make except-clause of copy() work
      * direct & inverse fields are both allocated upon initialization
      * Only rho is mutable.  To change other params, a new instance should be created
      * QUDA convention for clover field := 1+i ( kappa csw )/4 sigma_mu,nu F_mu,nu (<-sigma_mu,nu: spinor tensor)
      *  so that sigma_mu,nu = i[g_mu, g_nu], F_mu,nu = (Q_mu,nu - Q_nu,mu)/8 (1/2 is missing from sigma_mu,nu)
+     * Apparently, an input to QUDA clover object, coeff = kappa*csw
+     *  wihout a normalization factor of 1/4 or 1/32 (suggested in interface_quda.cpp)
     """
 
+    def __new__(cls, fmunu, **kwargs):
+        # TODO: get dofs and local dims from kwargs, instead of getting them
+        # from self.shape assuming that it has the form (dofs, local_dims)
+        if isinstance(fmunu, CloverField):
+            return fmunu
+        if not isinstance(fmunu, (numpy.ndarray, cupy.ndarray)):
+            raise TypeError(
+                f"Supporting only numpy or cupy for field, got {type(fmunu)}"
+            )
+        is_clover = False
+        if not isinstance(fmunu, GaugeField):
+            if kwargs.get("is_clover", False):
+                is_clover = True
+                field = fmunu
+            else:
+                fmunu = GaugeField(fmunu)
+
+        if not is_clover:  # not copying from a clover-field array
+            idof = int((fmunu.ncol * fmunu.ndims) ** 2 / 2)
+            prec = fmunu.dtype
+            field = fmunu.backend.empty((idof,) + fmunu.dims, dtype=prec)
+
+        return super().__new__(cls, field, **kwargs)
+
     def __init__(
         self,
-        fmunu,
+        obj,
         coeff=0.0,
         twisted=False,
         mu2=0,
         tf="NO",
         eps2=0,
         rho=0,
         computeTrLog=False,
+        **kwargs,
     ):
-        # ASSUME: coeff = kappa*csw wihout a normalization factor of 1/4 or 1/32 (suggested in interface_quda.cpp)
-        # ? better to store fmunu.quda_field to _fmunu -> import gauge_tensor to be used in some methods
-        # ? better to put clover into self.field -> need walk-around to make copy() work
-        if not isinstance(fmunu, GaugeField):
-            fmunu = GaugeField(fmunu)
-        self._fmunu = fmunu.compute_fmunu()
-        super().__init__(self._fmunu.field, comm=self._fmunu.comm)
-
-        # QUDA clover field inherently works with real's not with complex's (c.f., include/clover_field_order.h)
-
-        idof = int((self._fmunu.ncol * self._fmunu.ndims) ** 2 / 2)
-        prec = self._fmunu.precision
-        self._direct = (
-            False  # Here, it is a flag to indicate whether the field has been computed
-        )
-        self._inverse = (
-            False  # Here, it is a flag to indicate whether the field has been computed
-        )
-
-        new = lambda idof: LatticeField.create(
-            self._fmunu.global_lattice,
-            dofs=(idof,),
-            dtype=prec,
-            device=self._fmunu.device,
-            empty=True,
-        )
-        self._clover = new(idof)
-        self._cloverInv = new(idof)
-        self.coeff = coeff
-        self._twisted = twisted
-        self._twist_flavor = tf
-        self._mu2 = mu2
-        self._eps2 = eps2
-        self._rho = rho
-        self.computeTrLog = computeTrLog
-
-    # shape, dofs, dtype, iscomlex, isreal are overwriten to report their values for the clover field, instead of _fmunu
-
-    @property
-    def shape(self):
-        "Shape of the clover field"
-        return self._clover.shape
-
-    @property
-    def dofs(self):
-        "Shape of the per-site degrees of freedom"
-        return self._clover.dofs
-
-    @property
-    def dtype(self):
-        "Clover field data type"
-        return self._clover.dtype
-
-    @property
-    def iscomplex(self):
-        "Whether the clover field dtype is complex"
-        return self._clover.iscomplex
+        # WARNING: ndarray object is not supposed to be view-casted to CloverField object
+        #           except in __new__, for which __init__ will be called subsequently,
+        #           as the result won't come with  meta info such as 'coeff' or 'mu2'
+
+        super().__init__(obj, getattr(obj, "comm", None))
+        if isinstance(obj, GaugeField):
+            # explicit construction
+            # QUDA clover field inherently works with real's not with complex's (c.f., include/clover_field_order.h)
+            self._cloverInv = LatticeField.create(
+                self.global_lattice,
+                dofs=self.dofs,
+                dtype=self.dtype,
+                device=self.device_id,
+                empty=True,
+            )
+            self._fmunu = obj.compute_fmunu()
+            self._direct = False  # Here, it is a flag to indicate whether the field has been computed
+            self._inverse = False  # Here, it is a flag to indicate whether the field has been computed
+            self.coeff = coeff
+            self._twisted = twisted
+            self._twist_flavor = tf
+            self._mu2 = mu2
+            self._eps2 = eps2
+            self._rho = rho
+            self.computeTrLog = computeTrLog
+        elif isinstance(obj, CloverField):
+            # upcasting to ndarray or new from template
+            self.__dict__.update(obj.__dict__)
+        elif isinstance(obj, self.backend.ndarray):
+            pass
+        else:
+            raise ValueError(
+                "The input is expected to be ndarray or LatticeField object"
+            )
 
-    @property
-    def isreal(self):
-        "Whether the clover field dtype is real"
-        return self._clover.isreal
+    def _prepare(self, field, copy=False, check=False, **kwargs):
+        # When CloverField object prepares its input, the input is assumed to be of CloverField
+        return super()._prepare(field, copy=copy, check=check, is_clover=True, **kwargs)
 
     # naming suggestion: native_view? default_* saved for dofs+lattice?
     def default_view(self):
         N = 1 if self.order == "FLOAT2" else 4
         shape = (2,)  # even-odd
         shape += (self.dofs[0] // N, -1, N)
 
-        return self.field.view().reshape(shape)
+        return self.float_view().reshape(shape)
 
     @property
     def twisted(self):
@@ -178,7 +179,7 @@ def quda_params(self):
         params = lib.CloverFieldParam()
         lib.copy_struct(params, super().quda_params)
         params.inverse = True
-        params.clover = to_pointer(self._clover.ptr)
+        params.clover = to_pointer(self.ptr)
         params.cloverInv = to_pointer(self._cloverInv.ptr)
         params.coeff = self.coeff
         params.twisted = self.twisted
@@ -205,15 +206,15 @@ def clover_field(self):
         if not self._direct:
             lib.computeClover(self.quda_field, self._fmunu.quda_field, self.coeff)
             self._direct = True
-        return self._clover.field
+        return self.view(self.backend.ndarray)
 
     @property
     def inverse_field(self):
         if not self._inverse:
             self.clover_field
             lib.cloverInvert(self.quda_field, self.computeTrLog)
             self._inverse = True
-        return self._cloverInv.field
+        return self._cloverInv.view(self.backend.ndarray)
 
     @property
     def trLog(self):

lyncs_quda/dirac.py

@@ -242,9 +242,9 @@ def action(self, phi, **params):
         out = 0
         if not self.full and "CLOVER" in self.type:
             if self.even:
-                out -= 2 * self.clover.trLog[1]
-            else:
                 out -= 2 * self.clover.trLog[0]
+            else:
+                out -= 2 * self.clover.trLog[1]
 
         parity = None
         if not self.full:

lyncs_quda/gauge_field.py

@@ -74,11 +74,23 @@ def momentum(lattice, **kwargs):
 class GaugeField(LatticeField):
     "Mimics the quda::GaugeField object"
 
-    @LatticeField.field.setter
-    def field(self, field):
-        LatticeField.field.fset(self, field)
-        if self.reconstruct == "INVALID":
-            raise TypeError(f"Unrecognized field dofs {self.dofs}")
+    def _check_field(self, field=None):
+        # NOTE:
+        #  - For now, we store dofs and local lattice info in the shape of the array
+        #  - The canonical form for GaugeField is [dofs, local_lattice]
+        #     where len(dofs) == 1,2 and len(local_lattice) = 4
+        super()._check_field(field)
+        if field is None:
+            field = self
+        dofs = field.shape[: -self.ndims]
+        pdofs = (
+            prod(dofs)
+            if dofs[0] not in self._geometry_values[1] or dofs[0] == 1
+            else prod(dofs[1:])
+        )
+        pdofs *= 2 ** (self.iscomplex)
+        if not (pdofs in (12, 8, 10) or sqrt(pdofs / 2).is_integer()):
+            raise TypeError(f"Unrecognized field dofs {dofs}")
 
     def new(self, reconstruct=None, geometry=None, **kwargs):
         "Returns a new empty field based on the current"
@@ -139,11 +151,6 @@ def copy(self, other=None, out=None, **kwargs):
 
         return super().copy(other, out, **kwargs)
 
-    def __array_finalize__(self, obj):
-        "Support for __array_finalize__ standard"
-        # need to reset QUDA object when meta data of its Python wrapper is changed
-        self._quda = None
-
     def _prepare(self, field, **kwargs):
         field = super()._prepare(field, **kwargs)
         is_momentum = kwargs.get("is_momentum", self.is_momentum)
@@ -254,7 +261,7 @@ def is_momentum(self):
     def is_momentum(self, value):
         self._is_momentum = value
         if self._quda is not None:
-            self._quda.link_type = self.quda_link_type
+            self._quda = None
 
     @property
     def t_boundary(self):
@@ -523,7 +530,7 @@ def project(self, tol=None):
         if tol is None:
             tol = numpy.finfo(self.dtype).eps
 
-        assert self.device == cupy.cuda.runtime.getDevice()
+        assert self.device_id == cupy.cuda.runtime.getDevice()
         fails = cupy.zeros((1,), dtype="int32")
         lib.projectSU3(self.quda_field, tol, to_pointer(fails.data.ptr, "int *"))
         # return fails.get()[0]  # shouldn't we reduce?