test adding DataArray

This commit is contained in:
wyq 2025-12-03 14:49:53 +08:00
parent 153476bb93
commit d3f3fd894d
3 changed files with 58 additions and 9 deletions

View File

@ -1,15 +1,6 @@
from collections import OrderedDict
from pylib.mipylib import DimArray
class Coordinate(DimArray):
@property
def dim(self):
return self.dims[0]
class Coordinates(OrderedDict):

View File

@ -1,5 +1,7 @@
import mipylib.numeric as np
from .variable import Variable
from ._coordinates import Coordinates
class DataArray(np.NDArray):
@ -79,6 +81,37 @@ class DataArray(np.NDArray):
self._name = name
self._attrs = attrs
def array_wrap(self, arr, axis=None):
"""
Return a new array wrapped as self class object.
:param arr: The array to be wrapped.
:param axis: (*int*) The axis for ufunc compute along. Default is `None`, means not consider.
:return: New array object.
"""
if isinstance(arr, (Array, NDArray)):
if axis is None:
return DataArray(arr, self._coords, self._dims, self._name, self._attrs)
else:
dims = []
coords = Coordinates()
for i in range(0, self.ndim):
dim = self._dims[i]
if isinstance(axis, (list, tuple)):
if not i in axis:
dims.append(dim)
if dim in self._coords:
coords[dim] = self._coords[dim]
else:
if i != axis:
dims.append(dim)
if dim in self._coords:
coords[dim] = self._coords[dim]
return DataArray(arr, coords, dims, self._name, self._attrs)
else:
return arr
@property
def name(self):
"""The name of this array."""

View File

@ -10,6 +10,31 @@ class Variable(np.NDArray):
self._data = data
self._attrs = attrs
def array_wrap(self, arr, axis=None):
"""
Return a new array wrapped as self class object.
:param arr: The array to be wrapped.
:param axis: (*int*) The axis for ufunc compute along. Default is `None`, means not consider.
:return: New array object.
"""
if isinstance(arr, (Array, NDArray)):
if axis is None:
return Variable(arr, self._dims)
else:
dims = []
for i in range(0, self.ndim):
if isinstance(axis, (list, tuple)):
if not i in axis:
dims.append(self._dims[i])
else:
if i != axis:
dims.append(self._dims[i])
return Variable(arr, dims)
else:
return arr
@property
def data(self):
return self._data