|
4 | 4 |
|
5 | 5 | import numpy as np |
6 | 6 |
|
7 | | -from desc.backend import jnp, tree_flatten, tree_leaves, tree_unflatten |
| 7 | +from desc.backend import ( |
| 8 | + jnp, |
| 9 | + tree_flatten, |
| 10 | + tree_leaves, |
| 11 | + tree_map, |
| 12 | + tree_structure, |
| 13 | + tree_unflatten, |
| 14 | +) |
8 | 15 | from desc.compute import data_index |
9 | 16 | from desc.compute.utils import _compute as compute_fun |
10 | 17 | from desc.compute.utils import _parse_parameterization, get_profiles, get_transforms |
11 | 18 | from desc.grid import QuadratureGrid |
12 | 19 | from desc.optimizable import OptimizableCollection |
13 | | -from desc.utils import errorif, getsource, jaxify, parse_argname_change, setdefault |
| 20 | +from desc.utils import ( |
| 21 | + broadcast_tree, |
| 22 | + errorif, |
| 23 | + getsource, |
| 24 | + jaxify, |
| 25 | + parse_argname_change, |
| 26 | + setdefault, |
| 27 | +) |
14 | 28 |
|
15 | 29 | from .linear_objectives import _FixedObjective |
16 | 30 | from .objective_funs import _Objective, collect_docs |
@@ -669,3 +683,337 @@ def compute(self, params, constants=None): |
669 | 683 | ) |
670 | 684 | f = self._fun_wrapped(data) |
671 | 685 | return f |
| 686 | + |
| 687 | + |
| 688 | +class DeflationOperator(_Objective): |
| 689 | + r"""Deflation wrapper to be added to or to wrap objective to find new solutions. |
| 690 | +
|
| 691 | + If DeflationOperator is created while passing in an objective, the cost will be M*f |
| 692 | + where f is the objective's computed value. |
| 693 | + If DeflationOperator is created without passing in an objective, the cost will be |
| 694 | + only M |
| 695 | +
|
| 696 | + Deflation is done on the passed-in list of parameters. This objective |
| 697 | + value will be large if the current state is close to one of the already-found |
| 698 | + states given by `things_to_deflate`, thus enabling new solutions to be found, |
| 699 | + and guarantees that old solutions are not found (as the objective increases |
| 700 | + without bound as an already-found solution is approached) |
| 701 | +
|
| 702 | + The deflation operator is defined as: |
| 703 | +
|
| 704 | + M(x;xₖ)=(||x−xₖ||₂)⁻ᵖ + σ |
| 705 | +
|
| 706 | + (if `deflation_type="power"`) |
| 707 | +
|
| 708 | + or |
| 709 | +
|
| 710 | + M(𝐱;𝐱₁*) = exp(1/||𝐱−𝐱₁*||₂) + σ |
| 711 | +
|
| 712 | + (if `deflation_type="exp"`) |
| 713 | +
|
| 714 | + where x is the state and xₖ the passed-in known state. |
| 715 | + If multiple known states are used for deflation, then M is computed |
| 716 | + for each deflated state, then either multipled or added together (depending |
| 717 | + on if `multiple_deflation_type="prod"` or `"sum"`) to form the final cost. |
| 718 | + If an objective was passed in, this will then be multiplied by that objective's |
| 719 | + compute. |
| 720 | +
|
| 721 | + Parameters |
| 722 | + ---------- |
| 723 | + thing : Optimizable |
| 724 | + Optimizable that will be optimized to satisfy the Objective. |
| 725 | + things_to_deflate: list containing elements of type {Optimizable, None} |
| 726 | + list of objects to use in deflation operator. Should be same type |
| 727 | + as thing. Can also contain None elements, in which case those will be ignored. |
| 728 | + The utility of allowing the None element and ignoring them is if one is using |
| 729 | + this objective in a loop with a pre-determined number of iterations and adding |
| 730 | + each result of the loop iterate to the things_to_deflate, it may trigger |
| 731 | + recompilation of the objective's compute and jac/grad functions each time, |
| 732 | + which is wasteful. You can instead pass in a list containing None elements |
| 733 | + padding the list out to the max length it will attain. In this way, no |
| 734 | + recompilations will be triggered, and the entire loop will be completed |
| 735 | + much more quickly. |
| 736 | + If all things_to_deflate are None, this objective has zero cost (if not |
| 737 | + wrapping another objective) or simply returns the wrapped objective's |
| 738 | + cost (if wrapping another objective) |
| 739 | + params_to_deflate_with : nested list of dicts, optional |
| 740 | + Dict keys are the names of parameters to deflate (str), and dict values are the |
| 741 | + indices to deflate with for each corresponding parameter (int array). |
| 742 | + Use True (False) instead of an int array to deflate all (none) of the indices |
| 743 | + for that parameter. |
| 744 | + Must have the same pytree structure as thing.params_dict. |
| 745 | + The default is to deflate all indices of all parameters. |
| 746 | + objective: _Objective, optional |
| 747 | + Objective to wrap with the DeflationOperator. If not None, the cost will |
| 748 | + be M(x;xₖ)f(x) where f(x) is the Objective's cost. If None, then the cost |
| 749 | + returned will be M(x;xₖ). The objective must accept only one optimizable |
| 750 | + thing, and it must be the same as the thing passed to the DeflationOperator |
| 751 | + sigma: float, optional |
| 752 | + shift parameter in deflation operator. |
| 753 | + power: float, optional |
| 754 | + power parameter in deflation operator, ignored if `deflation_type="exp"`. |
| 755 | + deflation_type: {"power","exp"} |
| 756 | + What type of deflation to use. If `"power"`, uses the form |
| 757 | + pioneered by Farrell where M(𝐱;𝐱₁*) = ||𝐱−𝐱₁*||⁻ᵖ₂ + σ |
| 758 | + while `"exp"` uses the form from Riley 2024, where |
| 759 | + M(𝐱;𝐱₁*) = exp(1/||𝐱−𝐱₁*||₂) + σ. Defaults to "power". |
| 760 | + multiple_deflation_type: {"prod","sum"} |
| 761 | + When deflating multiple states, how to reduce the individual deflation |
| 762 | + terms Mᵢ(𝐱;𝐱ᵢ*). `"prod"` will multiply each individual deflation term |
| 763 | + together, while `"sum"` will add each individual term. |
| 764 | + single_shift: bool, |
| 765 | + Whether to use a single shift or include the shift in each individual |
| 766 | + deflation term. i.e. whether to use M = σ + prod(||𝐱−𝐱_i*||⁻ᵖ₂) (if True) |
| 767 | + or to use M = prod( σ + ||𝐱−𝐱_i*||⁻ᵖ₂). Defaults to False. |
| 768 | +
|
| 769 | + """ |
| 770 | + |
| 771 | + __doc__ = __doc__.rstrip() + collect_docs( |
| 772 | + target_default="``target=0``.", bounds_default="``target=0``." |
| 773 | + ) |
| 774 | + _static_attrs = _Objective._static_attrs + [ |
| 775 | + "_deflation_type", |
| 776 | + "_multiple_deflation_type", |
| 777 | + "_single_shift", |
| 778 | + "_params_to_deflate_with", |
| 779 | + ] |
| 780 | + |
| 781 | + _coordinates = "rtz" |
| 782 | + _units = "~" |
| 783 | + _print_value_fmt = "Deflation error: " |
| 784 | + |
| 785 | + def __init__( |
| 786 | + self, |
| 787 | + thing, |
| 788 | + things_to_deflate, |
| 789 | + params_to_deflate_with=None, |
| 790 | + objective=None, |
| 791 | + sigma=1.0, |
| 792 | + power=2, |
| 793 | + target=None, |
| 794 | + bounds=None, |
| 795 | + weight=1, |
| 796 | + normalize=True, |
| 797 | + normalize_target=True, |
| 798 | + loss_function=None, |
| 799 | + deriv_mode="auto", |
| 800 | + name="Deflation", |
| 801 | + jac_chunk_size=None, |
| 802 | + deflation_type="power", |
| 803 | + multiple_deflation_type="prod", |
| 804 | + single_shift=False, |
| 805 | + ): |
| 806 | + if target is None and bounds is None: |
| 807 | + target = 0 |
| 808 | + errorif( |
| 809 | + not np.all( |
| 810 | + [ |
| 811 | + (isinstance(t, type(thing)) and t != thing) or t is None |
| 812 | + for t in things_to_deflate |
| 813 | + ] |
| 814 | + ), |
| 815 | + ValueError, |
| 816 | + "All things_to_deflate must be the same type as" |
| 817 | + " thing and not the same object as thing.", |
| 818 | + ) |
| 819 | + self._things_to_deflate = things_to_deflate.copy() |
| 820 | + self._sigma = sigma |
| 821 | + self._power = power |
| 822 | + self._params_to_deflate_with = params_to_deflate_with |
| 823 | + errorif( |
| 824 | + deflation_type not in ["power", "exp"], |
| 825 | + ValueError, |
| 826 | + f"deflation_type must be 'power' or 'exp', got {deflation_type}", |
| 827 | + ) |
| 828 | + self._deflation_type = deflation_type |
| 829 | + errorif( |
| 830 | + multiple_deflation_type not in ["prod", "sum"], |
| 831 | + ValueError, |
| 832 | + "multiple_deflation_type must be 'prod' or 'sum'," |
| 833 | + f"got {multiple_deflation_type}", |
| 834 | + ) |
| 835 | + self._multiple_deflation_type = multiple_deflation_type |
| 836 | + self._single_shift = single_shift |
| 837 | + self._objective = objective |
| 838 | + if self._objective is not None: |
| 839 | + errorif( |
| 840 | + not isinstance(self._objective, _Objective), |
| 841 | + ValueError, |
| 842 | + "objective passed in must be an _Objective!", |
| 843 | + ) |
| 844 | + errorif( |
| 845 | + len(objective.things) > 1, |
| 846 | + NotImplementedError, |
| 847 | + "objective wrapped by DeflationOperator currently must have only" |
| 848 | + " a single object being optimized. Deflation on multiple optimizable" |
| 849 | + "objects at once is not yet implemented", |
| 850 | + ) |
| 851 | + errorif( |
| 852 | + objective.things[0] != thing, |
| 853 | + ValueError, |
| 854 | + "optimizable thing " |
| 855 | + " passed to DeflationOperator must be the same as the one passed to the" |
| 856 | + " wrapped objective", |
| 857 | + ) |
| 858 | + name = "Deflated " + self._objective._name |
| 859 | + self._units = self._objective._units |
| 860 | + self._scalar = self._objective._scalar |
| 861 | + self._coordinates = self._objective._coordinates |
| 862 | + self._print_value_fmt = "Deflated " + self._objective._print_value_fmt |
| 863 | + |
| 864 | + super().__init__( |
| 865 | + things=thing, |
| 866 | + target=target, |
| 867 | + bounds=bounds, |
| 868 | + weight=weight, |
| 869 | + normalize=normalize, |
| 870 | + normalize_target=normalize_target, |
| 871 | + loss_function=loss_function, |
| 872 | + deriv_mode=deriv_mode, |
| 873 | + name=name, |
| 874 | + jac_chunk_size=jac_chunk_size, |
| 875 | + ) |
| 876 | + |
| 877 | + def build(self, use_jit=True, verbose=1): |
| 878 | + """Build constant arrays. |
| 879 | +
|
| 880 | + Parameters |
| 881 | + ---------- |
| 882 | + use_jit : bool, optional |
| 883 | + Whether to just-in-time compile the objective and derivatives. |
| 884 | + verbose : int, optional |
| 885 | + Level of output. |
| 886 | +
|
| 887 | + """ |
| 888 | + thing = self.things[0] |
| 889 | + |
| 890 | + # default params |
| 891 | + default_params = tree_map(lambda dim: np.arange(dim), thing.dimensions) |
| 892 | + self._params_to_deflate_with = setdefault( |
| 893 | + self._params_to_deflate_with, default_params |
| 894 | + ) |
| 895 | + self._params_to_deflate_with = broadcast_tree( |
| 896 | + self._params_to_deflate_with, default_params |
| 897 | + ) |
| 898 | + self._indices = tree_leaves(self._params_to_deflate_with) |
| 899 | + errorif( |
| 900 | + tree_structure(self._params_to_deflate_with) |
| 901 | + != tree_structure(default_params), |
| 902 | + AssertionError, |
| 903 | + "", |
| 904 | + ) |
| 905 | + |
| 906 | + if self._objective is not None: |
| 907 | + if not self._objective.built: |
| 908 | + self._objective.build() |
| 909 | + self._dim_f = self._objective._dim_f |
| 910 | + self._normalization = self._objective._normalization |
| 911 | + self._constants = self._objective._constants |
| 912 | + else: |
| 913 | + self._dim_f = 1 |
| 914 | + |
| 915 | + self._is_not_none_mask = [] |
| 916 | + self._not_all_things_to_deflate_are_None = not np.all( |
| 917 | + [t is None for t in self._things_to_deflate] |
| 918 | + ) |
| 919 | + |
| 920 | + for i, t in enumerate(self._things_to_deflate): |
| 921 | + if t is None: |
| 922 | + self._is_not_none_mask.append(0.0) |
| 923 | + self._things_to_deflate[i] = thing |
| 924 | + else: |
| 925 | + self._is_not_none_mask.append(1.0) |
| 926 | + |
| 927 | + self._is_not_none_mask = np.array(self._is_not_none_mask, dtype=bool) |
| 928 | + |
| 929 | + if ( |
| 930 | + self._objective is None and self._bounds is not None |
| 931 | + ): # if being used as constraint/obj, min value should be sigma |
| 932 | + lower_bound_min = ( |
| 933 | + self._sigma |
| 934 | + if self._single_shift |
| 935 | + else self._sigma * np.sum(self._is_not_none_mask) |
| 936 | + ) |
| 937 | + errorif( |
| 938 | + not np.all(self._bounds[0] <= lower_bound_min), |
| 939 | + ValueError, |
| 940 | + ( |
| 941 | + f"Provided lower bound {self._bounds[0]} for deflation operator " |
| 942 | + f"is too high compared to the minimum value of {lower_bound_min} " |
| 943 | + "it can take based off of sigma, use a smaller lower bound" |
| 944 | + ), |
| 945 | + ) |
| 946 | + |
| 947 | + super().build(use_jit=use_jit, verbose=verbose) |
| 948 | + |
| 949 | + def compute(self, params, constants=None): |
| 950 | + """Compute deflation error. |
| 951 | +
|
| 952 | + Parameters |
| 953 | + ---------- |
| 954 | + params : dict |
| 955 | + Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict |
| 956 | + constants : dict |
| 957 | + Dictionary of constant data, eg transforms, profiles etc. Defaults to |
| 958 | + self.constants |
| 959 | +
|
| 960 | + Returns |
| 961 | + ------- |
| 962 | + f : scalar |
| 963 | + Deflation error. |
| 964 | +
|
| 965 | + """ |
| 966 | + this_thing_params = jnp.concatenate( |
| 967 | + [ |
| 968 | + jnp.atleast_1d(param[idx]) |
| 969 | + for param, idx in zip(tree_leaves(params), self._indices) |
| 970 | + ] |
| 971 | + ) |
| 972 | + diffs = [ |
| 973 | + this_thing_params |
| 974 | + - self._is_not_none_mask[i] |
| 975 | + * jnp.concatenate( |
| 976 | + [ |
| 977 | + jnp.atleast_1d(param[idx]) |
| 978 | + for param, idx in zip(tree_leaves(t.params_dict), self._indices) |
| 979 | + ] |
| 980 | + ) |
| 981 | + for i, t in enumerate(self._things_to_deflate) |
| 982 | + ] |
| 983 | + # to avoid division by zero if the states are the exact same |
| 984 | + eps = 1e2 * jnp.finfo(diffs[0].dtype).eps |
| 985 | + diffs = jnp.vstack(diffs) |
| 986 | + if self._deflation_type == "power": |
| 987 | + M_i = 1 / ( |
| 988 | + jnp.linalg.norm(diffs, axis=1) + eps |
| 989 | + ) ** self._power + self._sigma * (not self._single_shift) |
| 990 | + else: |
| 991 | + M_i = jnp.exp(1 / (jnp.linalg.norm(diffs, axis=1) + eps)) + self._sigma * ( |
| 992 | + not self._single_shift |
| 993 | + ) |
| 994 | + |
| 995 | + # we use the where= to only count the non-None things in things_to_deflate |
| 996 | + if self._multiple_deflation_type == "prod": |
| 997 | + deflation_parameter = jnp.prod( |
| 998 | + M_i, initial=1.0, where=self._is_not_none_mask |
| 999 | + ) + self._sigma * (self._single_shift) |
| 1000 | + else: |
| 1001 | + deflation_parameter = jnp.sum( |
| 1002 | + M_i, where=self._is_not_none_mask, initial=0.0 |
| 1003 | + ) + self._sigma * (self._single_shift) |
| 1004 | + |
| 1005 | + # enforce deflation paremeter 0 here if every thing_to_deflate is None |
| 1006 | + deflation_parameter *= self._not_all_things_to_deflate_are_None |
| 1007 | + |
| 1008 | + if self._objective is not None: |
| 1009 | + f = self._objective.compute(params) |
| 1010 | + # if wrapping an objective, but all things are None, make deflation do |
| 1011 | + # nothing when multiplying f, so here we add 1 to it as it is 0 right now |
| 1012 | + deflation_parameter += jnp.invert( |
| 1013 | + self._not_all_things_to_deflate_are_None |
| 1014 | + ).astype(float) |
| 1015 | + |
| 1016 | + else: |
| 1017 | + f = 1.0 |
| 1018 | + |
| 1019 | + return deflation_parameter * f |
0 commit comments