diff --git a/CHANGELOG.md b/CHANGELOG.md index 139f8ba85d..575803e029 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,8 @@ Breaking Changes and Deprecations Bug Fixes - Fixes bug in ``auglag`` optimizers which prevented them from accepting solver hyperparameters. +- Improves planar coil representation (``desc.coils.FourierPlanarCoil`` and ``desc.coils.FourierXYCoil``) internal rotation methods to avoid potential NaNs which could occur when the normal is parallel or antiparallel to Z-axis( to within machine epsilon), and also ensure the gradient at those edge cases is not only not NaN but also non-zero to avoid optimizer stalls at those cases. +- Fixes potential scaling-based issue in DESC-based optimization methods which use adaptive Hessian scaling (e.g. ``"fmintr"``) that could occur when the problem size was small and there were directions of near-zero derivative in the problem. - Fixes bug in modified Cholesky factorization used by the trust-region subproblems when the Gershgorin lower bound of the Hessian was exactly zero (e.g. a Hessian with an all-zero row), producing NaN steps in ``fmintr`` and diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index ef2c5b3aca..bf283386eb 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -5,11 +5,9 @@ from ..utils import ( cross, dot, - rotation_matrix, + rotate_vector_to_vector, rpz2xyz, rpz2xyz_vec, - safearccos, - safenormalize, xyz2rpz, xyz2rpz_vec, ) @@ -222,14 +220,7 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([X, Y, Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) coords = jnp.matmul(coords, A.T) + center coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"] # convert back to rpz @@ -267,14 +258,8 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([dX, dY, dZ]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) # convert back to rpz @@ -317,14 +302,8 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([d2X, d2Y, d2Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) # convert back to rpz @@ -374,14 +353,8 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([d3X, d3Y, d3Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) # convert back to rpz @@ -422,14 +395,8 @@ def _x_FourierXYCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([X, Y, Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) + center coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"] # convert back to rpz @@ -465,14 +432,8 @@ def _x_s_FourierXYCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([dX, dY, dZ]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) # convert back to rpz @@ -508,14 +469,8 @@ def _x_ss_FourierXYCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([d2X, d2Y, d2Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) # convert back to rpz @@ -551,14 +506,8 @@ def _x_sss_FourierXYCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([d3X, d3Y, d3Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, normal) - dotprod = dot(Zaxis, safenormalize(normal)) - angle = safearccos(dotprod) - A = jnp.where( # handle the case where normal is aligned with the -Z axis - jnp.allclose(dotprod, -1.0), - jnp.diag(jnp.array([1.0, -1.0, -1.0])), - rotation_matrix(axis, angle), - ) + A = rotate_vector_to_vector(Zaxis, normal) + coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) # convert back to rpz diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 38c70f764b..0d414c3cd8 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -190,7 +190,7 @@ def translate(self, displacement=[0, 0, 0]): def rotate(self, axis=[0, 0, 1], angle=0): """Rotate the curve by a fixed angle about axis in X,Y,Z coordinates.""" - R = rotation_matrix(axis=axis, angle=angle) + R = rotation_matrix(axis=np.asarray(axis).astype(float), angle=angle) self.rotmat = (R @ self.rotmat.reshape(3, 3)).flatten() self.shift = self.shift @ R.T diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index 54532ff336..79b9f36056 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -767,7 +767,7 @@ def normal(self): @normal.setter def normal(self, new): if len(np.asarray(new)) == 3: - self._normal = np.asarray(new) / np.linalg.norm(new) + self._normal = np.asarray(new).astype(float) / np.linalg.norm(new) else: raise ValueError( "normal should be a 3 element vector in " @@ -1121,7 +1121,7 @@ def normal(self): @normal.setter def normal(self, new): if len(np.asarray(new)) == 3: - self._normal = np.asarray(new) / np.linalg.norm(new) + self._normal = np.asarray(new).astype(float) / np.linalg.norm(new) else: raise ValueError( "normal should be a 3 element vector in " diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index d3b938eae0..80493bc308 100644 --- a/desc/optimize/utils.py +++ b/desc/optimize/utils.py @@ -518,7 +518,7 @@ def compute_hess_scale(H, prev_scale_inv=None): """Compute scaling factors based on diagonal of Hessian matrix.""" scale_inv = jnp.abs(jnp.diag(H)) scale_inv = jnp.where( - scale_inv < jnp.finfo(H.dtype).eps * max(H.shape), 1, scale_inv + scale_inv < jnp.finfo(H.dtype).eps * max(H.shape) * 1e1, 1, scale_inv ) if prev_scale_inv is not None: diff --git a/desc/utils.py b/desc/utils.py index 77f8627dfc..854400bc7f 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -983,6 +983,15 @@ def reflection_matrix(normal): def rotation_matrix(axis, angle=None): """Matrix to rotate points about axis by given angle. + NOTE: This function works but will have zero gradient w.r.t. + the angle when the angle is nearly zero (specifically, when + the angle