Skip to content

Commit b53635d

Browse files
committed
[WIP] Add support for nd vectors on md field
This is a first suggestion of how to implement spatio-temporal vector fields. It is hardly tested and more a foundation for a discussion.
1 parent 51204f4 commit b53635d

File tree

3 files changed

+24
-11
lines changed

3 files changed

+24
-11
lines changed

src/gstools/field/base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -578,7 +578,7 @@ def pos(self, pos):
578578
self._pos, self._field_shape = format_struct_pos_dim(pos, self.dim)
579579
# prepend dimension if we have a vector field
580580
if self.value_type == "vector":
581-
self._field_shape = (self.dim,) + self._field_shape
581+
self._field_shape = (self.generator.vec_dim,) + self._field_shape
582582
if self.latlon:
583583
raise ValueError("Field: Vector fields not allowed for latlon")
584584

src/gstools/field/generator.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,8 @@ class IncomprRandMeth(RandMeth):
339339
the mean velocity in x-direction
340340
mode_no : :class:`int`, optional
341341
number of Fourier modes. Default: ``1000``
342+
vec_dim : :class:`int`, optional
343+
vector dimension, in case it mismatches the model dimension
342344
seed : :class:`int` or :any:`None`, optional
343345
the seed of the random number generator.
344346
If "None", a random seed is used. Default: :any:`None`
@@ -391,18 +393,27 @@ def __init__(
391393
model,
392394
mean_velocity=1.0,
393395
mode_no=1000,
396+
vec_dim=None,
394397
seed=None,
395398
verbose=False,
396399
sampling="auto",
397400
**kwargs,
398401
):
399-
if model.dim < 2 or model.dim > 3:
402+
if vec_dim is None and (model.dim < 2 or model.dim > 3):
400403
raise ValueError(
401-
"Only 2D and 3D incompressible fields can be generated."
404+
"Only 2D and 3D incompressible vectors can be generated."
405+
)
406+
if vec_dim is not None and (vec_dim < 2 or vec_dim > 3):
407+
raise ValueError(
408+
"Only 2D and 3D incompressible vectors can be generated."
402409
)
403410
super().__init__(model, mode_no, seed, verbose, sampling, **kwargs)
404411

405412
self.mean_u = mean_velocity
413+
if vec_dim is None:
414+
self.vec_dim = model.dim
415+
else:
416+
self.vec_dim = vec_dim
406417
self._value_type = "vector"
407418

408419
def __call__(self, pos):
@@ -425,11 +436,12 @@ def __call__(self, pos):
425436
"""
426437
pos = np.asarray(pos, dtype=np.double)
427438
summed_modes = summate_incompr(
428-
self._cov_sample, self._z_1, self._z_2, pos
439+
self.vec_dim, self._cov_sample, self._z_1, self._z_2, pos
429440
)
430441
nugget = self.get_nugget(summed_modes.shape)
431442

432443
e1 = self._create_unit_vector(summed_modes.shape)
444+
extra_dim = max(0, self.model.dim - self.vec_dim)
433445

434446
return (
435447
self.mean_u * e1
@@ -458,7 +470,7 @@ def _create_unit_vector(self, broadcast_shape, axis=0):
458470
the unit vector
459471
"""
460472
shape = np.ones(len(broadcast_shape), dtype=int)
461-
shape[0] = self.model.dim
473+
shape[0] = self.vec_dim
462474

463475
e1 = np.zeros(shape)
464476
e1[axis] = 1.0

src/gstools/field/summator.pyx

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ cdef (double) abs_square(const double[:] vec) nogil:
5050

5151

5252
def summate_incompr(
53+
const int vec_dim,
5354
const double[:, :] cov_samples,
5455
const double[:] z_1,
5556
const double[:] z_2,
@@ -58,24 +59,24 @@ def summate_incompr(
5859
cdef int i, j, d
5960
cdef double phase
6061
cdef double k_2
61-
cdef int dim = pos.shape[0]
62+
cdef int field_dim = pos.shape[0]
6263

63-
cdef double[:] e1 = np.zeros(dim, dtype=float)
64+
cdef double[:] e1 = np.zeros(vec_dim, dtype=float)
6465
e1[0] = 1.
65-
cdef double[:] proj = np.empty(dim)
66+
cdef double[:] proj = np.empty(vec_dim)
6667

6768
cdef int X_len = pos.shape[1]
6869
cdef int N = cov_samples.shape[1]
6970

70-
cdef double[:, :] summed_modes = np.zeros((dim, X_len), dtype=float)
71+
cdef double[:, :] summed_modes = np.zeros((vec_dim, X_len), dtype=float)
7172

7273
for i in range(X_len):
7374
for j in range(N):
7475
k_2 = abs_square(cov_samples[:, j])
7576
phase = 0.
76-
for d in range(dim):
77+
for d in range(field_dim):
7778
phase += cov_samples[d, j] * pos[d, i]
78-
for d in range(dim):
79+
for d in range(vec_dim):
7980
proj[d] = e1[d] - cov_samples[d, j] * cov_samples[0, j] / k_2
8081
summed_modes[d, i] += proj[d] * (z_1[j] * cos(phase) + z_2[j] * sin(phase))
8182

0 commit comments

Comments
 (0)