Extract common function spaces from an array of forms. If ``forms`` is a list of linear forms, this function returns of list of the corresponding test function spaces. If ``forms`` is a 2D array of bilinear forms, for ``index=0`` the list of common test function spaces for each row
(
forms: Form | Sequence[Form] | Sequence[Sequence[Form]],
index: int = 0,
)
| 477 | |
| 478 | |
| 479 | def extract_function_spaces( |
| 480 | forms: Form | Sequence[Form] | Sequence[Sequence[Form]], |
| 481 | index: int = 0, |
| 482 | ) -> FunctionSpace | list[None | FunctionSpace]: |
| 483 | """Extract common function spaces from an array of forms. |
| 484 | |
| 485 | If ``forms`` is a list of linear forms, this function returns of list |
| 486 | of the corresponding test function spaces. If ``forms`` is a 2D |
| 487 | array of bilinear forms, for ``index=0`` the list of common test |
| 488 | function spaces for each row is returned, and if ``index=1`` the |
| 489 | common trial function spaces for each column are returned. |
| 490 | |
| 491 | Args: |
| 492 | forms: A list of forms or a 2D array of forms. |
| 493 | index: Index of the function space to extract. If ``index=0``, |
| 494 | the test function spaces are extracted, if ``index=1`` the |
| 495 | trial function spaces are extracted. |
| 496 | |
| 497 | Returns: |
| 498 | List of function spaces. |
| 499 | """ |
| 500 | _forms = np.array(forms) |
| 501 | if _forms.ndim == 0: |
| 502 | form: Form = _forms.tolist() |
| 503 | return form.function_spaces[0] if form is not None else None |
| 504 | elif _forms.ndim == 1: |
| 505 | assert index == 0, "Expected index=0 for 1D array of forms" |
| 506 | for form in _forms: |
| 507 | if form is not None: |
| 508 | assert form.rank == 1, "Expected linear form" |
| 509 | return [form.function_spaces[0] if form is not None else None for form in forms] # type: ignore[union-attr] |
| 510 | elif _forms.ndim == 2: |
| 511 | assert index == 0 or index == 1, "Expected index=0 or index=1 for 2D array of forms" |
| 512 | extract_spaces = np.vectorize( |
| 513 | lambda form: form.function_spaces[index] if form is not None else None |
| 514 | ) |
| 515 | V = extract_spaces(_forms) |
| 516 | |
| 517 | def unique_spaces(V): |
| 518 | # Pick spaces from first column |
| 519 | V0 = V[:, 0] |
| 520 | |
| 521 | # Iterate over each column |
| 522 | for col in range(1, V.shape[1]): |
| 523 | # Iterate over entry in column, updating if current |
| 524 | # space is None, or where both spaces are not None check |
| 525 | # that they are the same |
| 526 | for row in range(V.shape[0]): |
| 527 | if V0[row] is None and V[row, col] is not None: |
| 528 | V0[row] = V[row, col] |
| 529 | elif V0[row] is not None and V[row, col] is not None: |
| 530 | assert V0[row] is V[row, col], "Cannot extract unique function spaces" |
| 531 | return V0 |
| 532 | |
| 533 | if index == 0: |
| 534 | return list(unique_spaces(V)) |
| 535 | elif index == 1: |
| 536 | return list(unique_spaces(V.transpose())) |