Commit 05f7bdde authored by dboe's avatar dboe
Browse files

coord_sys was changing with empty grid. Fixed

parent e92ca9ba
Pipeline #106670 passed with stages
in 58 seconds
...@@ -57,10 +57,22 @@ class TensorGrid_Test(unittest.TestCase, TensorFields_Check): ...@@ -57,10 +57,22 @@ class TensorGrid_Test(unittest.TestCase, TensorFields_Check):
tge = tg.explicit() tge = tg.explicit()
self.check_filled(tge) self.check_filled(tge)
def test_coord_sys(self):
tg = TensorGrid.empty(
*self.bv, iter_order=self.iter_order, coord_sys=tfields.bases.CYLINDER
)
self.assertEqual(tg.coord_sys, tfields.bases.CYLINDER)
tge = tg.explicit()
self.assertEqual(tge.coord_sys, tfields.bases.CYLINDER)
def test_getitem(self): def test_getitem(self):
tg = TensorGrid.empty(*self.bv, iter_order=self.iter_order) tg = TensorGrid.empty(*self.bv, iter_order=self.iter_order)
self.check_filled(tg[:]) self.check_filled(tg[:])
self.check_filled(tg.explicit()[:]) tge = tg.explicit()
self.check_filled(tge[:])
tf = tge[:-2]
self.assertTrue(tf.equal(self.res[:-2]))
self.assertIsInstance(tf, tfields.TensorFields)
class TensorGrid_Test_Permutation1(TensorGrid_Test): class TensorGrid_Test_Permutation1(TensorGrid_Test):
......
"""
Implementaiton of TensorGrid class
"""
import numpy as np import numpy as np
from .lib import grid from .lib import grid
from .core import TensorFields from .core import TensorFields
...@@ -39,20 +42,40 @@ class TensorGrid(TensorFields): ...@@ -39,20 +42,40 @@ class TensorGrid(TensorFields):
def __getitem__(self, index): def __getitem__(self, index):
if not self.is_empty(): if not self.is_empty():
return super().__getitem__(index) return super().__getitem__(index)
return self.explicit().__getitem__(index) item = self.explicit()
if isinstance(index, slice):
if (
index.step in (None, 1)
and index.start in (None, 0)
and index.stop in (None, -1, len(item))
):
# full slice -> no type change
pass
else:
# downgrade
item = TensorFields(item)
return item.__getitem__(index)
@classmethod @classmethod
def from_base_vectors(cls, *base_vectors, tensors=None, fields=None, **kwargs): def from_base_vectors(cls, *base_vectors, tensors=None, fields=None, **kwargs):
"""
Build the grid (explicitly) from base vectors
"""
iter_order = kwargs.pop("iter_order", np.arange(len(base_vectors))) iter_order = kwargs.pop("iter_order", np.arange(len(base_vectors)))
if tensors is None: if tensors is None:
tensors = TensorFields.grid(*base_vectors, iter_order=iter_order, **kwargs) tensors = TensorFields.grid(*base_vectors, iter_order=iter_order, **kwargs)
obj = cls(tensors, base_vectors=base_vectors, iter_order=iter_order) obj = cls(tensors, base_vectors=base_vectors, iter_order=iter_order, **kwargs)
if fields: if fields:
# pylint:disable=attribute-defined-outside-init
obj.fields = fields obj.fields = fields
return obj return obj
@classmethod @classmethod
def empty(cls, *base_vectors, **kwargs): def empty(cls, *base_vectors, **kwargs):
"""
Build the grid (implicitly) from base vectors
"""
base_vectors = grid.ensure_complex(*base_vectors) base_vectors = grid.ensure_complex(*base_vectors)
bv_lengths = [int(bv[2].imag) for bv in base_vectors] bv_lengths = [int(bv[2].imag) for bv in base_vectors]
tensors = np.empty(shape=(np.prod(bv_lengths), 0)) tensors = np.empty(shape=(np.prod(bv_lengths), 0))
...@@ -68,7 +91,7 @@ class TensorGrid(TensorFields): ...@@ -68,7 +91,7 @@ class TensorGrid(TensorFields):
base_vectors = obj.base_vectors base_vectors = obj.base_vectors
else: else:
if not all( if not all(
[a == b for a, b in zip(base_vectors, obj.base_vectors)] (a == b for a, b in zip(base_vectors, obj.base_vectors))
): ):
raise NotImplementedError("Non-alligned base vectors") raise NotImplementedError("Non-alligned base vectors")
kwargs.setdefault("base_vectors", base_vectors) kwargs.setdefault("base_vectors", base_vectors)
...@@ -82,6 +105,10 @@ class TensorGrid(TensorFields): ...@@ -82,6 +105,10 @@ class TensorGrid(TensorFields):
return super().rank return super().rank
def is_empty(self): def is_empty(self):
"""
Check if the object is an implicit grid (base points are empty but base_vectors and iter
order can be used to build the explicit grid's base points).
"""
return 0 in self.shape return 0 in self.shape
def explicit(self): def explicit(self):
...@@ -93,11 +120,18 @@ class TensorGrid(TensorFields): ...@@ -93,11 +120,18 @@ class TensorGrid(TensorFields):
return self.from_base_vectors(*base_vectors, **kwargs) return self.from_base_vectors(*base_vectors, **kwargs)
def change_iter_order(self, iter_order): def change_iter_order(self, iter_order):
"""
Transform the iter order
"""
bv_lengths = [int(bv[2].imag) for bv in self.base_vectors] bv_lengths = [int(bv[2].imag) for bv in self.base_vectors]
field_swap_indices = grid.change_iter_order( field_swap_indices = grid.change_iter_order(
bv_lengths, self.iter_order, iter_order # pylint:disable=access-member-before-definition
bv_lengths,
self.iter_order,
iter_order,
) )
for field in self.fields: for field in self.fields:
field[:] = field[field_swap_indices] field[:] = field[field_swap_indices]
# pylint:disable=attribute-defined-outside-init
self.iter_order = iter_order self.iter_order = iter_order
self[:] = self.explicit() self[:] = self.explicit()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment