Coverage for src/km3pipe/math.py: 81%

1#!usr/bin/env python

2# -*- coding: utf-8 -*-

3# Filename: math.py

4# pylint: disable=C0103

5"""

6Maths, Geometry, coordinates.

7"""

8import numpy as np

10from .logger import get_logger

12__author__ = "Tamas Gal and Moritz Lotze"

14__credits__ = ["Vladimir Kulikovskiy"]

15__license__ = "MIT"

16__maintainer__ = "Tamas Gal and Moritz Lotze"

17__email__ = "tgal@km3net.de"

18__status__ = "Development"

20log = get_logger(__name__) # pylint: disable=C0103

23def neutrino_to_source_direction(phi, theta, radian=True):

24 """Flip the direction.

26 Parameters

27 ----------

28 phi, theta: neutrino direction

29 radian: bool [default=True]

30 receive + return angles in radian? (if false, use degree)

32 """

33 phi = np.atleast_1d(phi).copy()

34 theta = np.atleast_1d(theta).copy()

35 if not radian:

36 phi *= np.pi / 180

37 theta *= np.pi / 180

38 assert np.all(phi <= 2 * np.pi)

39 assert np.all(theta <= np.pi)

40 azimuth = (phi + np.pi) % (2 * np.pi)

41 zenith = np.pi - theta

42 if not radian:

43 azimuth *= 180 / np.pi

44 zenith *= 180 / np.pi

45 return azimuth, zenith

48def source_to_neutrino_direction(azimuth, zenith, radian=True):

49 """Flip the direction.

51 Parameters

52 ----------

53 zenith : float

54 neutrino origin

55 azimuth: float

56 neutrino origin

57 radian: bool [default=True]

58 receive + return angles in radian? (if false, use degree)

60 """

61 azimuth = np.atleast_1d(azimuth).copy()

62 zenith = np.atleast_1d(zenith).copy()

63 if not radian:

64 azimuth *= np.pi / 180

65 zenith *= np.pi / 180

66 phi = (azimuth - np.pi) % (2 * np.pi)

67 theta = np.pi - zenith

68 if not radian:

69 phi *= 180 / np.pi

70 theta *= 180 / np.pi

71 return phi, theta

74def theta(v):

75 """Neutrino direction in polar coordinates.

77 Parameters

78 ----------

79 v : array (x, y, z)

81 Notes

82 -----

83 Downgoing event: theta = 180deg

84 Horizont: 90deg

85 Upgoing: theta = 0

87 Angles in radians.

89 """

90 v = np.atleast_2d(v)

91 dir_z = v[:, 2]

92 return theta_separg(dir_z)

95def theta_separg(dir_z):

96 return np.arccos(dir_z)

99def phi(v):

100 """Neutrino direction in polar coordinates.

101

102 Parameters

103 ----------

104 v : array (x, y, z)

105

106

107 Notes

108 -----

109 ``phi``, ``theta`` is the opposite of ``zenith``, ``azimuth``.

110

111 Angles in radians.

112

113 """

114 v = np.atleast_2d(v)

115 dir_x = v[:, 0]

116 dir_y = v[:, 1]

117 return phi_separg(dir_x, dir_y)

118

119

120def phi_separg(dir_x, dir_y):

121 p = np.arctan2(dir_y, dir_x)

122 p[p < 0] += 2 * np.pi

123 return p

124

125

126def zenith(v):

127 """Return the zenith angle in radians.

128

129 Parameters

130 ----------

131 v : array (x, y, z)

132

133

134 Notes

135 -----

136 Defined as 'Angle respective to downgoing'.

137 Downgoing event: zenith = 0

138 Horizont: 90deg

139 Upgoing: zenith = 180deg

140

141 """

142 return angle_between((0, 0, -1), v)

143

144

145def azimuth(v):

146 """Return the azimuth angle in radians.

147

148 Parameters

149 ----------

150 v : array (x, y, z)

151

152

153 Notes

154 -----

155 ``phi``, ``theta`` is the opposite of ``zenith``, ``azimuth``.

156

157 This is the 'normal' azimuth definition -- beware of how you

158 define your coordinates. KM3NeT defines azimuth

159 differently than e.g. SLALIB, astropy, the AAS.org

160

161 """

162 v = np.atleast_2d(v)

163 azi = phi(v) - np.pi

164 azi[azi < 0] += 2 * np.pi

165 if len(azi) == 1:

166 return azi[0]

167 return azi

168

169

170def cartesian(phi, theta, radius=1):

171 x = radius * np.sin(theta) * np.cos(phi)

172 y = radius * np.sin(theta) * np.sin(phi)

173 z = radius * np.cos(theta)

174 return np.column_stack((x, y, z))

175

176

177def angle(v1, v2, normalized=False):

178 """

179 Compute the unsigned angle between two vectors. For a stacked input, the

180 angle is computed pairwise. Inspired by the "vg" Python package.

181

182 Parameters

183 ----------

184 v1 : np.array

185 With shape `(3,)` or a `kx3` stack of vectors.

186 v2 : np.array

187 A vector or stack of vectors with the same shape as `v1`.

188 normalized : bool

189 By default, the vectors will be normalised unless `normalized` is `True`.

190

191 Returns

192 -------

193 A float or a vector of floats with the angle in radians.

194

195 """

196 dot_products = np.einsum("ij,ij->i", v1.reshape(-1, 3), v2.reshape(-1, 3))

197

198 if normalized:

199 cosines = dot_products

200 else:

201 cosines = dot_products / magnitude(v1) / magnitude(v2)

202

203 # The dot product can exceed 1 or -1 and arccos will fail unless we clip

204 angles = np.arccos(np.clip(cosines, -1.0, 1.0))

205

206 if v1.ndim == v2.ndim == 1:

207 return angles[0]

208

209 return angles

210

211

212def magnitude(v):

213 """

214 Calculates the magnitude of a vector or array of vectors.

215

216 Parameters

217 ----------

218 v : np.array

219 With shape `(3,)` or `kx3` stack of vectors.

220

221 Returns

222 -------

223 A float or a vector of floats with the magnitudes.

224 """

225 if v.ndim == 1:

226 return np.linalg.norm(v)

227 elif v.ndim == 2:

228 return np.linalg.norm(v, axis=1)

229 else:

230 ValueError("Unsupported dimensions")

231

232

233def angle_between(v1, v2, axis=0):

234 """Returns the angle in radians between vectors 'v1' and 'v2'.

235

236 If axis=1 it evaluates the angle between arrays of vectors.

237

238 Examples

239 --------

240 >>> angle_between((1, 0, 0), (0, 1, 0))

241 1.5707963267948966

242 >>> angle_between((1, 0, 0), (1, 0, 0))

243 0.0

244 >>> angle_between((1, 0, 0), (-1, 0, 0))

245 3.141592653589793

246 >>> v1 = np.array([[1, 0, 0], [1, 0, 0]])

247 >>> v2 = np.array([[0, 1, 0], [-1, 0, 0]])

248 >>> angle_between(v1, v2, axis=1)

249 array([1.57079633, 3.14159265])

250

251 """

252 if axis == 0:

253 v1_u = unit_vector(v1)

254 v2_u = unit_vector(v2)

255 # Don't use `np.dot`, does not work with all shapes

256 return np.arccos(np.clip(np.inner(v1_u, v2_u), -1, 1))

257 elif axis == 1:

258 return angle(v1, v2) # returns angle in deg

259 else:

260 raise ValueError("unsupported axis")

261

262

263def unit_vector(vector, **kwargs):

264 """Returns the unit vector of the vector."""

265 vector = np.array(vector)

266 out_shape = vector.shape

267 vector = np.atleast_2d(vector)

268 unit = vector / np.linalg.norm(vector, axis=1, **kwargs)[:, None]

269 return unit.reshape(out_shape)

270

271

272def pld3(pos, line_vertex, line_dir):

273 """Calculate the point-line-distance for given point and line."""

274 pos = np.atleast_2d(pos)

275 line_vertex = np.atleast_1d(line_vertex)

276 line_dir = np.atleast_1d(line_dir)

277 c = np.cross(line_dir, line_vertex - pos)

278 n1 = np.linalg.norm(c, axis=1)

279 n2 = np.linalg.norm(line_dir)

280 out = n1 / n2

281 if out.ndim == 1 and len(out) == 1:

282 return out[0]

283 return out

284

285

286def lpnorm(x, p=2):

287 return np.power(np.sum(np.power(x, p)), 1 / p)

288

289

290def dist(x1, x2, axis=0):

291 """Return the distance between two points.

292

293 Set axis=1 if x1 is a vector and x2 a matrix to get a vector of distances.

294 """

295 return np.linalg.norm(x2 - x1, axis=axis)

296

297

298def com(points, masses=None):

299 """Calculate center of mass for given points.

300 If masses is not set, assume equal masses."""

301 if masses is None:

302 return np.average(points, axis=0)

303 else:

304 return np.average(points, axis=0, weights=masses)

305

306

307def circ_permutation(items):

308 """Calculate the circular permutation for a given list of items."""

309 permutations = []

310 for i in range(len(items)):

311 permutations.append(items[i:] + items[:i])

312 return permutations

313

314

315def hsin(theta):

316 """haversine"""

317 return (1.0 - np.cos(theta)) / 2.0

318

319

320def space_angle(phi1, theta1, phi2, theta2):

321 """Also called Great-circle-distance --

322

323 use long-ass formula from wikipedia (last in section):

324 https://en.wikipedia.org/wiki/Great-circle_distance#Computational_formulas

325

326 Space angle only makes sense in lon-lat, so convert zenith -> latitude.

327

328 """

329 from numpy import pi, sin, cos, arctan2, sqrt, square

330

331 lamb1 = pi / 2 - theta1

332 lamb2 = pi / 2 - theta2

333 lambdelt = lamb2 - lamb1

334 under = sin(phi1) * sin(phi2) + cos(phi1) * cos(phi2) * cos(lambdelt)

335 over = sqrt(

336 np.square((cos(phi2) * sin(lambdelt)))

337 + square(cos(phi1) * sin(phi2) - sin(phi1) * cos(phi2) * cos(lambdelt))

338 )

339 angle = arctan2(over, under)

340 return angle

341

342

343def rotation_matrix(axis, theta):

344 """The Euler–Rodrigues formula.

345

346 Return the rotation matrix associated with counterclockwise rotation about

347 the given axis by theta radians.

348

349 Parameters

350 ----------

351 axis: vector to rotate around

352 theta: rotation angle [rad]

353

354 """

355 axis = np.asarray(axis)

356 axis = axis / np.linalg.norm(axis)

357 a = np.cos(theta / 2)

358 b, c, d = -axis * np.sin(theta / 2)

359 aa, bb, cc, dd = a * a, b * b, c * c, d * d

360 bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d

361 return np.array(

362 [

363 [aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],

364 [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],

365 [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc],

366 ]

367 )

368

369

370def spherecutmask(center, rmin, rmax, items):

371 """Returns a mask to select items, within a certain radius around a given center.

372

373

374 Parameters

375 -----------

376 center: array of shape(3,)

377 central point of the sphere

378 rmin: float

379 minimum radius of the sphere selection

380 rmax: float

381 maximum radius of the sphere selection

382 items: array of shape (n, 3) or iterable with pos_[xyz]-attributed items

383 the items to cut on

384

385 Returns

386 --------

387 mask of the array of shape (n, 3) or iterable with pos_[xyz]-attributed items.

388 """

389 items_pos = items

390

391 if all(hasattr(items, "pos_" + q) for q in "xyz"):

392 items_pos = np.array([items.pos_x, items.pos_y, items.pos_z]).T

393

394 distances = dist(center, items_pos, axis=1)

395 mask = (distances >= rmin) & (distances <= rmax)

396

397 return mask

398

399

400def spherecut(center, rmin, rmax, items):

401 """Select items within a certain radius around a given center.

402 This function calls spherecutmask() to create the selection mask and returns the selected items.

403

404 Parameters

405 -----------

406 center: array of shape(3,)

407 central point of the sphere

408 rmin: float

409 minimum radius of the sphere selection

410 rmax: float

411 maximum radius of the sphere selection

412 items: array of shape (n, 3) or iterable with pos_[xyz]-attributed items

413 the items to cut on

414

415 Returns

416 --------

417 array of shape (k, 3) or iterable with pos_[xyz]-attributed items

418 items which survived the cut.

419

420 """

421 selected_items = items[spherecutmask(center, rmin, rmax, items)]

422

423 return selected_items

424

425

426class Polygon(object):

427 """A polygon, to implement containment conditions."""

428

429 def __init__(self, vertices):

430 from matplotlib.path import Path

431

432 self.poly = Path(vertices)

433

434 def contains(self, points):

435 points = np.atleast_2d(points)

436 points_flat = points.reshape((-1, 2))

437 is_contained = self.poly.contains_points(points_flat)

438 return is_contained

439

440 def contains_xy(self, x, y):

441 xy = np.column_stack((x, y))

442 return self.contains(xy)

443

444

445class IrregularPrism(object):

446 """Like a cylinder, but the top is an irregular Polygon."""

447

448 def __init__(self, xy_vertices, z_min, z_max):

449 self.poly = Polygon(xy_vertices)

450 self.z_min = z_min

451 self.z_max = z_max

452

453 def _is_z_contained(self, z):

454 return (self.z_min <= z) & (z <= self.z_max)

455

456 def contains(self, points):

457 points = np.atleast_2d(points)

458 points_xy = points[:, [0, 1]]

459 points_z = points[:, 2]

460 is_xy_contained = self.poly.contains(points_xy)

461 is_z_contained = self._is_z_contained(points_z)

462 return is_xy_contained & is_z_contained

463

464 def contains_xyz(self, x, y, z):

465 xyz = np.column_stack((x, y, z))

466 return self.contains(xyz)

467

468

469class SparseCone(object):

470 """A Cone, represented by sparse samples.

471

472 This samples evenly spaced points from the base circle.

473

474 Parameters

475 ----------

476 spike_pos : coordinates of the top

477 bottom_center_pos : center of the bottom circle

478 opening_angle : cone opening angle, in rad

479 theta, axis to mantle, *not* mantle-mantle. So this is the angle

480 to the axis, and mantle-to-mantle (aperture) is 2 theta.

481 """

482

483 def __init__(self, spike_pos, bottom_center_pos, opening_angle):

484 self.spike_pos = np.asarray(spike_pos)

485 self.bottom_center_pos = np.asarray(bottom_center_pos)

486 self.opening_angle = opening_angle

487 self.top_bottom_vec = self.bottom_center_pos - self.spike_pos

488 self.height = np.linalg.norm(self.top_bottom_vec)

489 self.mantle_length = self.height / np.cos(self.opening_angle)

490 self.radius = self.height * np.tan(self.opening_angle * self.height)

491

492 @classmethod

493 def _equidistant_angles_from_circle(cls, n_angles=4):

494 return np.linspace(0, 2 * np.pi, n_angles + 1)[:-1]

495

496 @property

497 def _random_circle_vector(self):

498 k = self.top_bottom_vec

499 r = self.radius

500 x = np.random.randn(3)

501 x -= x.dot(k) * k

502 x *= r / np.linalg.norm(x)

503 return x

504

505 def sample_circle(self, n_angles=4):

506 angles = self._equidistant_angles_from_circle(n_angles)

507 random_circle_vector = self._random_circle_vector

508 # rotate the radius vector around the cone axis

509 points_on_circle = [

510 np.dot(rotation_matrix(self.top_bottom_vec, theta), random_circle_vector)

511 for theta in angles

512 ]

513 return points_on_circle

514

515 @property

516 def sample_axis(self):

517 return [self.spike_pos, self.bottom_center_pos]

518

519 def sample(self, n_circle=4):

520 points = self.sample_circle(n_circle)

521 points.extend(self.sample_axis)

522 return points

523

524

525def inertia(x, y, z, weight=None):

526 """Inertia tensor, stolen of thomas"""

527 if weight is None:

528 weight = 1

529 tensor_of_inertia = np.zeros((3, 3), dtype=float)

530 tensor_of_inertia[0][0] = (y * y + z * z) * weight

531 tensor_of_inertia[0][1] = (-1) * x * y * weight

532 tensor_of_inertia[0][2] = (-1) * x * z * weight

533 tensor_of_inertia[1][0] = (-1) * x * y * weight

534 tensor_of_inertia[1][1] = (x * x + z * z) * weight

535 tensor_of_inertia[1][2] = (-1) * y * z * weight

536 tensor_of_inertia[2][0] = (-1) * x * z * weight

537 tensor_of_inertia[2][1] = (-1) * z * y * weight

538 tensor_of_inertia[2][2] = (x * x + y * y) * weight

539

540 eigen_values = np.linalg.eigvals(tensor_of_inertia)

541 small_inertia = eigen_values[2][2]

542 middle_inertia = eigen_values[1][1]

543 big_inertia = eigen_values[0][0]

544 return small_inertia, middle_inertia, big_inertia

545

546

547def g_parameter(time_residual):

548 """stolen from thomas"""

549 mean = np.mean(time_residual)

550 time_residual_prime = time_residual - np.ones(time_residual.shape) * mean

551 time_residual_prime *= time_residual_prime / (-2 * 1.5 * 1.5)

552 time_residual_prime = np.exp(time_residual_prime)

553 g = np.sum(time_residual_prime) / len(time_residual)

554 return g

555

556

557def gold_parameter(time_residual):

558 """stolen from thomas"""

559 gold = np.exp(-1 * time_residual * time_residual / (2 * 1.5 * 1.5)) / len(

560 time_residual

561 )

562 gold = np.sum(gold)

563

564

565def log_b(arg, base):

566 """Logarithm to any base"""

567 return np.log(arg) / np.log(base)

568

569

570def qrot(vector, quaternion):

571 """Rotate a 3D vector using quaternion algebra.

572

573 Implemented by Vladimir Kulikovskiy.

574

575 Parameters

576 ----------

577 vector : np.array

578 quaternion : np.array

579

580 Returns

581 -------

582 np.array

583

584 """

585 t = 2 * np.cross(quaternion[1:], vector)

586 v_rot = vector + quaternion[0] * t + np.cross(quaternion[1:], t)

587 return v_rot

588

589

590def qeuler(yaw, pitch, roll):

591 """Convert Euler angle to quaternion.

592

593 Parameters

594 ----------

595 yaw : number

596 pitch : number

597 roll : number

598

599 Returns

600 -------

601 np.array

602

603 """

604 yaw = np.radians(yaw)

605 pitch = np.radians(pitch)

606 roll = np.radians(roll)

607

608 cy = np.cos(yaw * 0.5)

609 sy = np.sin(yaw * 0.5)

610 cr = np.cos(roll * 0.5)

611 sr = np.sin(roll * 0.5)

612 cp = np.cos(pitch * 0.5)

613 sp = np.sin(pitch * 0.5)

614

615 q = np.array(

616 (

617 cy * cr * cp + sy * sr * sp,

618 cy * sr * cp - sy * cr * sp,

619 cy * cr * sp + sy * sr * cp,

620 sy * cr * cp - cy * sr * sp,

621 )

622 )

623 return q

624

625

626def qrot_yaw(vector, heading):

627 """Rotate vectors using quaternion algebra.

628

629

630 Parameters

631 ----------

632 vector : np.array or list-like (3 elements)

633 heading : the heading to rotate to [deg]

634

635 Returns

636 -------

637 np.array

638

639 """

640 return qrot(vector, qeuler(heading, 0, 0))

641

642

643def intersect_3d(p1, p2):

644 """Find the closes point for a given set of lines in 3D.

645

646 Parameters

647 ----------

648 p1 : (M, N) array_like

649 Starting points

650 p2 : (M, N) array_like

651 End points.

652

653 Returns

654 -------

655 x : (N,) ndarray

656 Least-squares solution - the closest point of the intersections.

657

658 Raises

659 ------

660 numpy.linalg.LinAlgError

661 If computation does not converge.

662

663 """

664 v = p2 - p1

665 normed_v = unit_vector(v)

666 nx = normed_v[:, 0]

667 ny = normed_v[:, 1]

668 nz = normed_v[:, 2]

669 xx = np.sum(nx**2 - 1)

670 yy = np.sum(ny**2 - 1)

671 zz = np.sum(nz**2 - 1)

672 xy = np.sum(nx * ny)

673 xz = np.sum(nx * nz)

674 yz = np.sum(ny * nz)

675 M = np.array([(xx, xy, xz), (xy, yy, yz), (xz, yz, zz)])

676 x = np.sum(p1[:, 0] * (nx**2 - 1) + p1[:, 1] * (nx * ny) + p1[:, 2] * (nx * nz))

677 y = np.sum(p1[:, 0] * (nx * ny) + p1[:, 1] * (ny * ny - 1) + p1[:, 2] * (ny * nz))

678 z = np.sum(p1[:, 0] * (nx * nz) + p1[:, 1] * (ny * nz) + p1[:, 2] * (nz**2 - 1))

679 return np.linalg.lstsq(M, np.array((x, y, z)), rcond=None)[0]