From 3e5ff4e6725498a344f6af1fe4a412e59f55ccf3 Mon Sep 17 00:00:00 2001 From: Eivind Fonn Date: Tue, 13 Feb 2024 15:53:03 +0100 Subject: [PATCH] Mypy type safety: round 4 --- splipy/splineobject.py | 1 + splipy/volume_factory.py | 213 +++++++++++++++++++++++++-------------- 2 files changed, 141 insertions(+), 73 deletions(-) diff --git a/splipy/splineobject.py b/splipy/splineobject.py index f1ae2d9..0e8c10e 100644 --- a/splipy/splineobject.py +++ b/splipy/splineobject.py @@ -65,6 +65,7 @@ class SplineObject: _intended_pardim: ClassVar[Optional[int]] = None + dimension: int bases: list[BSplineBasis] controlpoints: FArray rational: bool diff --git a/splipy/volume_factory.py b/splipy/volume_factory.py index 35371f9..c48c3ff 100644 --- a/splipy/volume_factory.py +++ b/splipy/volume_factory.py @@ -3,22 +3,30 @@ """Handy utilities for creating volumes.""" from math import pi, sqrt, atan2 +from typing import Union, Literal, overload, Sequence, cast, Optional +from itertools import chain, repeat import numpy as np from .basis import BSplineBasis +from .curve import Curve from .surface import Surface from .volume import Volume from .utils import flip_and_move_plane_geometry, rotate_local_x_axis +from .utils.curve import curve_length_parametrization +from .types import Scalar, Scalars, FArray from . import curve_factory, surface_factory + __all__ = ['cube', 'sphere', 'revolve', 'cylinder', 'extrude', 'edge_surfaces', 'loft', 'interpolate', 'least_square_fit'] - -def cube(size=1, lower_left=(0,0,0)): - """ Create a cube with parmetric origin at *(0,0,0)*. +def cube( + size: Union[Scalar, tuple[Scalar, Scalar, Scalar]] = 1, + lower_left: Scalars = (0,0,0), +) -> Volume: + """Create a cube with parmetric origin at *(0,0,0)*. :param float size: Size(s), either a single scalar or a tuple of scalars per axis :param array-like lower_left: local origin, the lower bottom left corner of the cube @@ -26,12 +34,20 @@ def cube(size=1, lower_left=(0,0,0)): :rtype: Volume """ result = Volume() - result.scale(size) + if isinstance(size, tuple): + result.scale(*size) + else: + result.scale(size) result += lower_left return result -def sphere(r=1, center=(0,0,0), type='radial'): - """ Create a solid sphere + +def sphere( + r: Scalar = 1, + center: Scalars = (0,0,0), + type: Literal["radial", "square"] = 'radial', +) -> Volume: + """Create a solid sphere :param float r: Radius :param array-like center: Local origin of the sphere @@ -39,10 +55,12 @@ def sphere(r=1, center=(0,0,0), type='radial'): :return: A solid ball :rtype: Volume """ + if type == 'radial': shell = surface_factory.sphere(r, center) midpoint = shell*0 + center return edge_surfaces(shell, midpoint) + elif type == 'square': # based on the work of James E.Cobb: "Tiling the Sphere with Rational Bezier Patches" # University of Utah, July 11, 1988. UUCS-88-009 @@ -50,38 +68,45 @@ def sphere(r=1, center=(0,0,0), type='radial'): sr2 = sqrt(2) sr3 = sqrt(3) sr6 = sqrt(6) - cp = [[ -4*(sr3-1), 4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], # row 0 - [ -sr2 , sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], - [ 0 , 4./3*(1-2*sr3), 4./3*(1-2*sr3),4./3*(5-sr3) ], - [ sr2 , sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], - [ 4*(sr3-1), 4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], - [ -sr2*(4-sr3), -sr2, sr2*(sr3-4), sr2*(3*sr3-2)], # row 1 - [ -(3*sr3-2)/2, (2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], - [ 0 , sr2*(2*sr3-7)/3, -5*sr6/3, sr2*(sr3+6)/3], - [ (3*sr3-2)/2, (2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], - [ sr2*(4-sr3), -sr2, sr2*(sr3-4), sr2*(3*sr3-2)], - [ -4./3*(2*sr3-1), 0, 4./3*(1-2*sr3), 4*(5-sr3)/3], # row 2 - [-sr2/3*(7-2*sr3), 0, -5*sr6/3, sr2*(sr3+6)/3], - [ 0 , 0, 4*(sr3-5)/3, 4*(5*sr3-1)/9], - [ sr2/3*(7-2*sr3), 0, -5*sr6/3, sr2*(sr3+6)/3], - [ 4./3*(2*sr3-1), 0, 4./3*(1-2*sr3), 4*(5-sr3)/3], - [ -sr2*(4-sr3), sr2, sr2*(sr3-4), sr2*(3*sr3-2)], # row 3 - [ -(3*sr3-2)/2, -(2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], - [ 0 ,-sr2*(2*sr3-7)/3, -5*sr6/3, sr2*(sr3+6)/3], - [ (3*sr3-2)/2, -(2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], - [ sr2*(4-sr3), sr2, sr2*(sr3-4), sr2*(3*sr3-2)], - [ -4*(sr3-1), -4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], # row 4 - [ -sr2 , -sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], - [ 0 , -4./3*(1-2*sr3), 4./3*(1-2*sr3),4./3*(5-sr3) ], - [ sr2 , -sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], - [ 4*(sr3-1), -4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ]] - wmin = Surface(b,b,cp, rational=True) + cp = np.array( + [ + [ -4*(sr3-1), 4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], # row 0 + [ -sr2 , sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], + [ 0 , 4./3*(1-2*sr3), 4./3*(1-2*sr3),4./3*(5-sr3) ], + [ sr2 , sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], + [ 4*(sr3-1), 4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], + [ -sr2*(4-sr3), -sr2, sr2*(sr3-4), sr2*(3*sr3-2)], # row 1 + [ -(3*sr3-2)/2, (2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], + [ 0 , sr2*(2*sr3-7)/3, -5*sr6/3, sr2*(sr3+6)/3], + [ (3*sr3-2)/2, (2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], + [ sr2*(4-sr3), -sr2, sr2*(sr3-4), sr2*(3*sr3-2)], + [ -4./3*(2*sr3-1), 0, 4./3*(1-2*sr3), 4*(5-sr3)/3], # row 2 + [-sr2/3*(7-2*sr3), 0, -5*sr6/3, sr2*(sr3+6)/3], + [ 0 , 0, 4*(sr3-5)/3, 4*(5*sr3-1)/9], + [ sr2/3*(7-2*sr3), 0, -5*sr6/3, sr2*(sr3+6)/3], + [ 4./3*(2*sr3-1), 0, 4./3*(1-2*sr3), 4*(5-sr3)/3], + [ -sr2*(4-sr3), sr2, sr2*(sr3-4), sr2*(3*sr3-2)], # row 3 + [ -(3*sr3-2)/2, -(2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], + [ 0 ,-sr2*(2*sr3-7)/3, -5*sr6/3, sr2*(sr3+6)/3], + [ (3*sr3-2)/2, -(2-3*sr3)/2, -(sr3+6)/2, (sr3+6)/2], + [ sr2*(4-sr3), sr2, sr2*(sr3-4), sr2*(3*sr3-2)], + [ -4*(sr3-1), -4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], # row 4 + [ -sr2 , -sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], + [ 0 , -4./3*(1-2*sr3), 4./3*(1-2*sr3),4./3*(5-sr3) ], + [ sr2 , -sr2*(sr3-4), sr2*(sr3-4), sr2*(3*sr3-2)], + [ 4*(sr3-1), -4*(1-sr3), 4*(1-sr3), 4*(3-sr3) ], + ], + dtype=float, + ) + + wmin = Surface(b, b, cp, rational=True) wmax = wmin.clone().mirror([0,0,1]) vmax = wmin.clone().rotate(pi/2, [1,0,0]) vmin = vmax.clone().mirror([0,1,0]) umax = vmin.clone().rotate(pi/2, [0,0,1]) umin = umax.clone().mirror([1,0,0]) - # ideally I would like to call edge_surfaces() now, but that function + + # Ideally I would like to call edge_surfaces() now, but that function # does not work with rational surfaces, so we'll just manually try # and add some inner controlpoints cp = np.zeros((5,5,5,4)) @@ -98,13 +123,18 @@ def sphere(r=1, center=(0,0,0), type='radial'): cp[1:4,1:4,1:4,2] = Z cp[1:4,1:4,1:4,3] = 1 ball = Volume(b,b,b,cp,rational=True, raw=True) - return r*ball + center + return r * ball + center + else: raise ValueError('invalid type argument') -def revolve(surf, theta=2 * pi, axis=(0,0,1)): - """ Revolve a volume by sweeping a surface in a rotational fashion around +def revolve( + surf: Surface, + theta: Scalar = 2 * pi, + axis: Scalars = (0,0,1), +) -> Volume: + """Revolve a volume by sweeping a surface in a rotational fashion around an axis. :param Surface surf: Surface to revolve @@ -142,8 +172,16 @@ def revolve(surf, theta=2 * pi, axis=(0,0,1)): result.rotate(normal_theta, [0,0,1]) return result -def torus(minor_r=1, major_r=3, center=(0,0,0), normal=(0,0,1), xaxis=(1,0,0), type='radial'): - """ Create a torus (doughnut) by revolving a circle of size *minor_r* + +def torus( + minor_r: Scalar = 1, + major_r: Scalar = 3, + center: Scalars = (0,0,0), + normal: Scalars = (0,0,1), + xaxis: Scalars = (1,0,0), + type: Literal["radial", "square"] = 'radial', +) -> Volume: + """Create a torus (doughnut) by revolving a circle of size *minor_r* around the *z* axis with radius *major_r*. :param float minor_r: The thickness of the torus (radius in the *xz* plane) @@ -162,10 +200,18 @@ def torus(minor_r=1, major_r=3, center=(0,0,0), normal=(0,0,1), xaxis=(1,0,0), t result = revolve(disc) result.rotate(rotate_local_x_axis(xaxis, normal)) - return flip_and_move_plane_geometry(result, center, normal) + return result.flip_and_move_plane_geometry(center, normal) + -def cylinder(r=1, h=1, center=(0,0,0), axis=(0,0,1), xaxis=(1,0,0), type='radial'): - """ Create a solid cylinder +def cylinder( + r: Scalar = 1, + h: Scalar = 1, + center: Scalars = (0,0,0), + axis: Scalars = (0,0,1), + xaxis: Scalars = (1,0,0), + type: Literal["radial", "square"] = 'radial', +) -> Volume: + """Create a solid cylinder :param float r: Radius :param float h: Height @@ -179,8 +225,8 @@ def cylinder(r=1, h=1, center=(0,0,0), axis=(0,0,1), xaxis=(1,0,0), type='radial return extrude(surface_factory.disc(r, center, axis, xaxis=xaxis, type=type), h*np.array(axis)) -def extrude(surf, amount): - """ Extrude a surface by sweeping it to a given height. +def extrude(surf: Surface, amount: Scalars) -> Volume: + """Extrude a surface by sweeping it to a given height. :param Surface surf: Surface to extrude :param array-like amount: 3-component vector of sweeping amount and direction @@ -188,18 +234,27 @@ def extrude(surf, amount): :rtype: Volume """ surf.set_dimension(3) # add z-components (if not already present) - cp = [] - for controlpoint in surf: - cp.append(list(controlpoint)) + + ncomps = surf.dimension + surf.rational + cp = surf.controlpoints.copy().reshape(-1, ncomps, order='F') surf += amount - for controlpoint in surf: - cp.append(list(controlpoint)) + cp = np.append(cp, surf.controlpoints.reshape(-1, ncomps, order='F'), axis=0) surf -= amount - return Volume(surf.bases[0], surf.bases[1], BSplineBasis(2), cp, surf.rational) + return Volume(surf.bases[0], surf.bases[1], BSplineBasis(2), cp, rational=surf.rational) + + +@overload +def edge_surfaces(surfaces: Sequence[Surface]) -> Volume: + ... -def edge_surfaces(*surfaces): - """ Create the volume defined by the region between the input surfaces. +@overload +def edge_surfaces(*surfaces: Surface) -> Volume: + ... + + +def edge_surfaces(*surfaces: Surface) -> Volume: # type: ignore[misc] + """Create the volume defined by the region between the input surfaces. In case of six input surfaces, these must be given in the order: bottom, top, left, right, back, front. Opposing sides must be parametrized in the @@ -211,7 +266,8 @@ def edge_surfaces(*surfaces): :raises ValueError: If the length of *surfaces* is not two or six """ if len(surfaces) == 1: # probably gives input as a list-like single variable - surfaces = surfaces[0] + surfaces = cast(tuple[Surface], surfaces[0]) + if len(surfaces) == 2: surf1 = surfaces[0].clone() surf2 = surfaces[1].clone() @@ -317,8 +373,9 @@ def edge_surfaces(*surfaces): else: raise ValueError('Requires two or six input surfaces') -def sweep(path, shape): - """ Generate a surface by sweeping a shape along a path + +def sweep(path: Curve, shape: Surface) -> Volume: + """Generate a surface by sweeping a shape along a path The resulting surface is an approximation generated by interpolating at the Greville points. It is generated by sweeping a shape curve along a path. @@ -358,8 +415,16 @@ def sweep(path, shape): return interpolate(X, [b1,b2,b3]) -def loft(*surfaces): - """ Generate a volume by lofting a series of surfaces +@overload +def loft(surfaces: Sequence[Surface]) -> Volume: + ... + +@overload +def loft(*surfaces: Surface) -> Volume: + ... + +def loft(*surfaces: Surface) -> Volume: # type: ignore[misc] + """Generate a volume by lofting a series of surfaces The resulting volume is interpolated at all input surfaces and a smooth transition between these surfaces is computed as a cubic spline interpolation in the lofting @@ -392,27 +457,29 @@ def loft(*surfaces): """ if len(surfaces) == 1: - surfaces = surfaces[0] + surfaces = cast(tuple[Surface], surfaces[0]) # clone input, so we don't change those references # make sure everything has the same dimension since we need to compute length - surfaces = [s.clone().set_dimension(3) for s in surfaces] + surfaces = tuple(s.clone().set_dimension(3) for s in surfaces) + if len(surfaces)==2: - return surface_factory.edge_curves(surfaces) - elif len(surfaces)==3: + return edge_surfaces(surfaces) + + if len(surfaces)==3: # can't do cubic spline interpolation, so we'll do quadratic basis3 = BSplineBasis(3) dist = basis3.greville() + else: x = [s.center() for s in surfaces] # create knot vector from the euclidian length between the surfaces - dist = [0] - for (x1,x0) in zip(x[1:],x[:-1]): - dist.append(dist[-1] + np.linalg.norm(x1-x0)) + dist = np.zeros((len(x),), dtype=float) + curve_length_parametrization(x, buffer=dist[1:]) # using "free" boundary condition by setting N'''(u) continuous at second to last and second knot - knot = [dist[0]]*4 + dist[2:-2] + [dist[-1]]*4 + knot = np.fromiter(chain(repeat(dist[0], 4), dist[2:-2], repeat(dist[-1], 4)), dtype=float) basis3 = BSplineBasis(4, knot) n = len(surfaces) @@ -438,24 +505,24 @@ def loft(*surfaces): Nw_inv = np.linalg.inv(Nw) # compute interpolation points in physical space - x = np.zeros((m1,m2,n, dim)) + xx = np.zeros((m1, m2, n, dim), dtype=float) for i in range(n): - tmp = np.tensordot(Nv, surfaces[i].controlpoints, axes=(1,1)) - x[:,:,i,:] = np.tensordot(Nu, tmp , axes=(1,1)) + tmp = np.tensordot(Nv, surfaces[i].controlpoints, axes=(1,1)) + xx[:,:,i,:] = np.tensordot(Nu, tmp, axes=(1,1)) # solve interpolation problem - cp = np.tensordot(Nw_inv, x, axes=(1,2)) + cp = np.tensordot(Nw_inv, xx, axes=(1,2)) cp = np.tensordot(Nv_inv, cp, axes=(1,2)) cp = np.tensordot(Nu_inv, cp, axes=(1,2)) # re-order controlpoints so they match up with Surface constructor cp = np.reshape(cp.transpose((2, 1, 0, 3)), (m1*m2*n, dim)) - return Volume(basis1, basis2, basis3, cp, surfaces[0].rational) + return Volume(basis1, basis2, basis3, cp, rational=surfaces[0].rational) -def interpolate(x, bases, u=None): - """ Interpolate a volume on a set of regular gridded interpolation points `x`. +def interpolate(x: FArray, bases: Sequence[BSplineBasis], u: Optional[Sequence[Scalars]] = None) -> Volume: + """Interpolate a volume on a set of regular gridded interpolation points `x`. The points can be either a matrix (in which case the first index is interpreted as a flat row-first index of the interpolation grid) or a 4D @@ -483,8 +550,8 @@ def interpolate(x, bases, u=None): return Volume(bases[0], bases[1], bases[2], cp.transpose(2,1,0,3).reshape((np.prod(vol_shape),dim))) -def least_square_fit(x, bases, u): - """ Perform a least-square fit of a point cloud `x` onto a spline basis. +def least_square_fit(x: FArray, bases: Sequence[BSplineBasis], u: Sequence[Scalars]) -> Volume: + """Perform a least-square fit of a point cloud `x` onto a spline basis. The points can be either a matrix (in which case the first index is interpreted as a flat row-first index of the interpolation grid) or a 4D