Skip to content

Commit

Permalink
Mypy type safety: round 4
Browse files Browse the repository at this point in the history
  • Loading branch information
@TheBB
TheBB committed Feb 13, 2024
1 parent 76322c9 commit 3e5ff4e
Show file tree
Hide file tree
Showing 2 changed files with 141 additions and 73 deletions.
1 change: 1 addition & 0 deletions splipy/splineobject.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ class SplineObject:

_intended_pardim: ClassVar[Optional[int]] = None

dimension: int
bases: list[BSplineBasis]
controlpoints: FArray
rational: bool
Expand Down
213 changes: 140 additions & 73 deletions splipy/volume_factory.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,85 +3,110 @@
"""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
:return: A linear parametrized box
: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
:param string type: The type of parametrization ('radial' or 'square')
: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
b = BSplineBasis(order=5)
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))
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -179,27 +225,36 @@ 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
:return: The extruded surface
: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
Expand All @@ -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()
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 3e5ff4e

Please sign in to comment.