|
| 1 | +import numpy as np |
| 2 | +from scipy.interpolate import BSpline |
| 3 | + |
| 4 | + |
| 5 | +class BSplineAiry: |
| 6 | + """ |
| 7 | + Manages a tensor-product B-spline surface for the Airy stress function. |
| 8 | +
|
| 9 | + This class pre-computes basis functions to allow fast evaluation of |
| 10 | + stress fields (derivatives of the Airy function) from a set of control points. |
| 11 | +
|
| 12 | + The Airy stress function phi(x,y) is represented as: |
| 13 | + phi(x,y) = sum_i sum_j C_ij * B_i(x) * B_j(y) |
| 14 | +
|
| 15 | + Stresses are: |
| 16 | + sigma_xx = d^2(phi)/dy^2 |
| 17 | + sigma_yy = d^2(phi)/dx^2 |
| 18 | + sigma_xy = -d^2(phi)/dxdy |
| 19 | + """ |
| 20 | + |
| 21 | + def __init__(self, shape, knot_spacing, degree=3): |
| 22 | + """ |
| 23 | + Initialize the B-spline basis for a given image shape. |
| 24 | +
|
| 25 | + Parameters |
| 26 | + ---------- |
| 27 | + shape : tuple |
| 28 | + Image shape (height, width). |
| 29 | + knot_spacing : int |
| 30 | + Approximate spacing between knots in pixels. |
| 31 | + degree : int |
| 32 | + Degree of the B-spline (default: 3 for cubic). |
| 33 | + """ |
| 34 | + self.ny, self.nx = shape |
| 35 | + self.degree = degree |
| 36 | + self.knot_spacing = knot_spacing |
| 37 | + |
| 38 | + # Generate knots |
| 39 | + # We need knots to cover the range [0, n] with sufficient padding for the degree |
| 40 | + # interior knots |
| 41 | + tx = np.arange(0, self.nx + knot_spacing, knot_spacing) |
| 42 | + ty = np.arange(0, self.ny + knot_spacing, knot_spacing) |
| 43 | + |
| 44 | + # Add simpler padding |
| 45 | + # Standard way for clamped B-spline is repeating start/end knots degree+1 times |
| 46 | + # But for general coverage we just need enough support. |
| 47 | + # Let's use the standard "clamped" knot vector construction for the domain [0, L] |
| 48 | + |
| 49 | + def make_knots(length, step): |
| 50 | + # Interior knots |
| 51 | + interior = np.arange(0, length + step, step) |
| 52 | + if interior[-1] < length: |
| 53 | + interior = np.append(interior, length) |
| 54 | + |
| 55 | + # Pad with clamped ends |
| 56 | + t = np.concatenate(([interior[0]] * degree, interior, [interior[-1]] * degree)) |
| 57 | + return t |
| 58 | + |
| 59 | + self.tx = make_knots(self.nx, knot_spacing) |
| 60 | + self.ty = make_knots(self.ny, knot_spacing) |
| 61 | + |
| 62 | + # Number of coefficients (control points) |
| 63 | + self.n_coeffs_x = len(self.tx) - degree - 1 |
| 64 | + self.n_coeffs_y = len(self.ty) - degree - 1 |
| 65 | + self.n_coeffs = self.n_coeffs_x * self.n_coeffs_y |
| 66 | + |
| 67 | + # Pre-evaluate basis functions on the pixel grid |
| 68 | + # We evaluate at pixel centers |
| 69 | + x_grid = np.arange(self.nx) |
| 70 | + y_grid = np.arange(self.ny) |
| 71 | + |
| 72 | + # Evaluate basis functions B(x) and derivatives |
| 73 | + # This creates matrices of shape (width, n_coeffs_x) |
| 74 | + self.Bx, self.dBx, self.ddBx = self._precompute_basis(x_grid, self.tx, self.n_coeffs_x) |
| 75 | + self.By, self.dBy, self.ddBy = self._precompute_basis(y_grid, self.ty, self.n_coeffs_y) |
| 76 | + |
| 77 | + def _precompute_basis(self, coords, knots, n_coeffs): |
| 78 | + """ |
| 79 | + Compute B-spline basis matrix and its 1st and 2nd derivatives. |
| 80 | + Returns matrices of shape (len(coords), n_coeffs). |
| 81 | + """ |
| 82 | + # We use scipy BSpline.design_matrix-like logic but explicit |
| 83 | + # We want to know the value of the i-th basis function at each coordinate. |
| 84 | + # B_mat[k, i] = B_i(coords[k]) |
| 85 | + |
| 86 | + B_mat = np.zeros((len(coords), n_coeffs)) |
| 87 | + dB_mat = np.zeros((len(coords), n_coeffs)) |
| 88 | + ddB_mat = np.zeros((len(coords), n_coeffs)) |
| 89 | + |
| 90 | + # Iterate over each basis function |
| 91 | + # This might be slow for very large grids, but it's done once. |
| 92 | + # A faster way is to realize only degree+1 functions are non-zero at any point. |
| 93 | + # But optimizing this pre-calc is secondary to the main loop speed. |
| 94 | + |
| 95 | + for i in range(n_coeffs): |
| 96 | + # Create a localized BSpline for the i-th basis function |
| 97 | + # The coefficient vector is 1 at i and 0 elsewhere |
| 98 | + c = np.zeros(n_coeffs) |
| 99 | + c[i] = 1.0 |
| 100 | + spl = BSpline(knots, c, self.degree) |
| 101 | + |
| 102 | + B_mat[:, i] = spl(coords) |
| 103 | + dB_mat[:, i] = spl(coords, nu=1) |
| 104 | + ddB_mat[:, i] = spl(coords, nu=2) |
| 105 | + |
| 106 | + return B_mat, dB_mat, ddB_mat |
| 107 | + |
| 108 | + def get_stress_fields(self, coeffs_flat): |
| 109 | + """ |
| 110 | + Compute stress fields from flat coefficient array. |
| 111 | +
|
| 112 | + Parameters |
| 113 | + ---------- |
| 114 | + coeffs_flat : array-like |
| 115 | + Flattened array of coefficients of length n_coeffs_x * n_coeffs_y. |
| 116 | +
|
| 117 | + Returns |
| 118 | + ------- |
| 119 | + sigma_xx, sigma_yy, sigma_xy : ndarray |
| 120 | + Stress fields of shape (height, width). |
| 121 | + """ |
| 122 | + C = coeffs_flat.reshape(self.n_coeffs_y, self.n_coeffs_x) |
| 123 | + |
| 124 | + # sigma_xx = d^2(phi)/dy^2 = By'' * C * Bx.T |
| 125 | + # Shape: (ny, n_cy) @ (n_cy, n_cx) @ (n_cx, nx) -> (ny, nx) |
| 126 | + sigma_xx = self.ddBy @ C @ self.Bx.T |
| 127 | + |
| 128 | + # sigma_yy = d^2(phi)/dx^2 = By * C * Bx''.T |
| 129 | + sigma_yy = self.By @ C @ self.ddBx.T |
| 130 | + |
| 131 | + # sigma_xy = -d^2(phi)/dxdy = -(By' * C * Bx'.T) |
| 132 | + # Note: In image coordinates (y down), d/dy_img = -d/y_phy. |
| 133 | + # The cross derivative term gains a negative sign from the coordinate flip, |
| 134 | + # cancelling the negative sign in the Airey definition. |
| 135 | + # So sigma_xy = + d^2(phi)/dx_img dy_img |
| 136 | + sigma_xy = self.dBy @ C @ self.dBx.T |
| 137 | + |
| 138 | + return sigma_xx, sigma_yy, sigma_xy |
0 commit comments