Compute the area spanned by the unit vectors in arr and ref. Args: arr: (*b_shape, *d_shape, 3) ref: (*b_shape, *d_shape, 3) ndim_b: number of dimension of b_shape. If None, = 0. normalized: whetehr arr and ref are unit vectors
(
arr: torch.Tensor,
ref: torch.Tensor,
ndim_b: int = None,
normalized: bool = True,
valid_mask: torch.Tensor = None,
)
| 343 | |
| 344 | |
| 345 | def compute_area( |
| 346 | arr: torch.Tensor, |
| 347 | ref: torch.Tensor, |
| 348 | ndim_b: int = None, |
| 349 | normalized: bool = True, |
| 350 | valid_mask: torch.Tensor = None, |
| 351 | ): |
| 352 | """ |
| 353 | Compute the area spanned by the unit vectors in arr and ref. |
| 354 | |
| 355 | Args: |
| 356 | arr: (*b_shape, *d_shape, 3) |
| 357 | ref: (*b_shape, *d_shape, 3) |
| 358 | ndim_b: |
| 359 | number of dimension of b_shape. If None, = 0. |
| 360 | normalized: |
| 361 | whetehr arr and ref are unit vectors |
| 362 | valid_mask: (*b_shape, *d_shape,) |
| 363 | |
| 364 | Returns: |
| 365 | err: (*b_shape,) |
| 366 | """ |
| 367 | if ndim_b is None: |
| 368 | ndim_b = 0 |
| 369 | |
| 370 | if not normalized: |
| 371 | arr = torch.nn.functional.normalize(arr, p=2, dim=-1) |
| 372 | ref = torch.nn.functional.normalize(ref, p=2, dim=-1) |
| 373 | |
| 374 | out = torch.linalg.cross(arr, ref, dim=-1) # (*b, *d, 3) |
| 375 | area = torch.linalg.vector_norm(out, ord=2, dim=-1) # (*b, *d,) |
| 376 | |
| 377 | if valid_mask is None: |
| 378 | area = area.reshape(*(arr.shape[:ndim_b]), -1) # (*b, numel_d) |
| 379 | area = area.mean(dim=-1) # (*b,) |
| 380 | else: |
| 381 | valid_mask = valid_mask.view( |
| 382 | *(valid_mask.shape), *([1] * (area.ndim - valid_mask.ndim))).expand_as(area) |
| 383 | valid_mask = valid_mask.reshape(*(arr.shape[:ndim_b]), -1) # (*b, numel_d) |
| 384 | area = area.reshape(*(arr.shape[:ndim_b]), -1) # (*b, numel_d) |
| 385 | area = (area * valid_mask).sum(dim=-1) / valid_mask.sum(-1) |
| 386 | |
| 387 | return area |
| 388 | |
| 389 | |
| 390 | def compute_diff_angle( |