diff --git a/nutils/function.py b/nutils/function.py index 261e4b876..6dfee58a6 100644 --- a/nutils/function.py +++ b/nutils/function.py @@ -177,6 +177,17 @@ def cast(cls, __value: IntoArray, dtype: Optional[DType] = None, ndim: Optional[ raise ValueError('expected an array with dimension `{}` but got `{}`'.format(ndim, value.ndim)) return value + @classmethod + def cast_withscale(cls, __value: IntoArray, dtype: Optional[DType] = None, ndim: Optional[int] = None): + try: + scale = type(__value).reference_quantity + except AttributeError: + value = cls.cast(__value, dtype=dtype, ndim=ndim) + scale = value.dtype(1) + else: + value = cls.cast(__value / scale, dtype=dtype, ndim=ndim) + return value, scale + def __init__(self, shape: Shape, dtype: DType, spaces: FrozenSet[str], arguments: Mapping[str, Tuple[Shape, DType]]) -> None: self.shape = tuple(sh.__index__() for sh in shape) self.dtype = dtype @@ -2732,7 +2743,7 @@ def derivative(__arg: IntoArray, __var: Union[str, 'Argument']) -> Array: :class:`Array` ''' - arg = Array.cast(__arg) + arg, argscale = Array.cast_withscale(__arg) if isinstance(__var, str): if __var not in arg.arguments: raise ValueError('no such argument: {}'.format(__var)) @@ -2746,7 +2757,7 @@ def derivative(__arg: IntoArray, __var: Union[str, 'Argument']) -> Array: raise ValueError('Argument {!r} has shape {} in the function, but the derivative to {!r} with shape {} was requested.'.format(__var.name, shape, __var.name, __var.shape)) if __var.dtype != dtype: raise ValueError('Argument {!r} has dtype {} in the function, but the derivative to {!r} with dtype {} was requested.'.format(__var.name, dtype.__name__ if dtype in _dtypes else dtype, __var.name, __var.dtype.__name__ if __var.dtype in _dtypes else __var.dtype)) - return _Derivative(arg, __var) + return _Derivative(arg, __var) * argscale def grad(__arg: IntoArray, __geom: IntoArray, ndims: int = 0) -> Array: '''Return the gradient of the argument to the given geometry. @@ -2762,21 +2773,22 @@ def grad(__arg: IntoArray, __geom: IntoArray, ndims: int = 0) -> Array: :class:`Array` ''' - arg = Array.cast(__arg) - geom = Array.cast(__geom) + arg, argscale = Array.cast_withscale(__arg) + geom, geomscale = Array.cast_withscale(__geom) if geom.dtype != float: raise ValueError('The geometry must be real-valued.') if geom.ndim == 0: - return grad(arg, _append_axes(geom, (1,)))[...,0] + ret = grad(arg, _append_axes(geom, (1,)))[...,0] elif geom.ndim > 1: sh = geom.shape[-2:] - return unravel(grad(arg, ravel(geom, geom.ndim-2)), arg.ndim+geom.ndim-2, sh) + ret = unravel(grad(arg, ravel(geom, geom.ndim-2)), arg.ndim+geom.ndim-2, sh) elif ndims == 0 or ndims == geom.shape[0]: - return _Gradient(arg, geom) + ret = _Gradient(arg, geom) elif ndims == -1 or ndims == geom.shape[0] - 1: - return _SurfaceGradient(arg, geom) + ret = _SurfaceGradient(arg, geom) else: raise NotImplementedError + return ret * (argscale / geomscale) def curl(__arg: IntoArray, __geom: IntoArray) -> Array: '''Return the curl of the argument w.r.t. the given geometry. @@ -2819,7 +2831,7 @@ def normal(__geom: IntoArray, refgeom: Optional[Array] = None) -> Array: :class:`Array` ''' - geom = Array.cast(__geom) + geom, geomscale_ = Array.cast_withscale(__geom) if geom.dtype != float: raise ValueError('The geometry must be real-valued.') if geom.ndim == 0: @@ -2888,15 +2900,22 @@ def jacobian(__geom: IntoArray, __ndims: Optional[int] = None) -> Array: :class:`Array` ''' - geom = Array.cast(__geom) + geom, geomscale = Array.cast_withscale(__geom) if geom.dtype != float: raise ValueError('The geometry must be real-valued.') + if __ndims is not None: + scale = geomscale**__ndims + elif _is_unit_scalar(geomscale): + scale = 1. + else: + raise ValueError('scaled arrays require an explicit dimension') if geom.ndim == 0: - return jacobian(insertaxis(geom, 0, 1), __ndims) + J = jacobian(insertaxis(geom, 0, 1), __ndims) elif geom.ndim > 1: - return jacobian(ravel(geom, geom.ndim-2), __ndims) + J = jacobian(ravel(geom, geom.ndim-2), __ndims) else: - return _Jacobian(geom, __ndims) + J = _Jacobian(geom, __ndims) + return J * scale def J(__geom: IntoArray, __ndims: Optional[int] = None) -> Array: '''Return the absolute value of the determinant of the Jacobian matrix of the given geometry. @@ -3044,7 +3063,8 @@ def eval(funcs: evaluable.AsEvaluableArray, **arguments: Mapping[str, numpy.ndar results : :class:`tuple` of arrays ''' - return map(sparse.toarray, evaluable.eval_sparse(funcs, **arguments)) + funcs, funcscales = zip(*map(Array.cast_withscale, funcs)) + return tuple(sparse.toarray(data) * scale for data, scale in zip(evaluable.eval_sparse(funcs, **arguments), funcscales)) def isarray(__arg: Any) -> bool: 'Test if the argument is an instance of :class:`Array`.' diff --git a/nutils/sample.py b/nutils/sample.py index 469762cd7..0e155a6b5 100644 --- a/nutils/sample.py +++ b/nutils/sample.py @@ -189,9 +189,10 @@ def integrate(self, funcs: Iterable[function.IntoArray], arguments: Mapping[str, Optional arguments for function evaluation. ''' + funcs, funcscales = zip(*map(function.Array.cast_withscale, funcs)) datas = self.integrate_sparse(funcs, argdict(arguments)) with log.iter.fraction('assembling', datas) as items: - return tuple(_convert(data, inplace=True) for data in items) + return tuple(_convert(data, inplace=True) * scale for data, scale in zip(items, funcscales)) @util.single_or_multiple def integrate_sparse(self, funcs: Iterable[function.IntoArray], arguments: Optional[Mapping[str, numpy.ndarray]] = None) -> Tuple[numpy.ndarray, ...]: @@ -216,7 +217,8 @@ def integral(self, __func: function.IntoArray) -> function.Array: Integrand. ''' - return _Integral(function.Array.cast(__func), self) + func, funcscale = function.Array.cast_withscale(__func) + return _Integral(func, self) * funcscale @util.positional_only @util.single_or_multiple @@ -231,9 +233,10 @@ def eval(self, funcs: Iterable[function.IntoArray], arguments: Mapping[str, nump Optional arguments for function evaluation. ''' + funcs, funcscales = zip(*map(function.Array.cast_withscale, funcs)) datas = self.eval_sparse(funcs, arguments) with log.iter.fraction('assembling', datas) as items: - return tuple(map(sparse.toarray, items)) + return tuple(sparse.toarray(data) * funcscale for data, funcscale in zip(datas, funcscales)) @util.positional_only @util.single_or_multiple @@ -251,10 +254,10 @@ def eval_sparse(self, funcs: Iterable[function.IntoArray], arguments: Optional[M return evaluable.eval_sparse(map(self, funcs), **(arguments or {})) def __call__(self, __func: function.IntoArray) -> function.Array: - func = _ConcatenatePoints(function.Array.cast(__func), self) + func, funcscale = function.Array.cast_withscale(__func) ielem = evaluable.loop_index('_sample_' + '_'.join(self.spaces), self.nelems) indices = evaluable.loop_concatenate(evaluable._flat(self.get_evaluable_indices(ielem)), ielem) - return _ReorderPoints(func, indices) + return _ReorderPoints(_ConcatenatePoints(func, self), indices) * funcscale def basis(self) -> function.Array: '''Basis-like function that for every point in the sample evaluates to the