torch_vn_algebra.optimize¶

Source: torch_vn_algebra/optimize.py. Signatures and docstrings as in the source.

Module documentation¶

Batched gradient-based optimisation over operators of a Type I algebra.

Monte Carlo sampling explores typical operators; extremal values of a functional (the
sharp constants in an inequality) sit on sets of tiny measure that random sampling almost never
hits in higher dimensions. This module parametrises the usual constraint sets differentiably,
so that sup / inf of a functional can be searched for with many independent starts at once
(one start per batch element):

    UnitaryParam            U = U0 exp(G - G^*)  (Haar-random U0 per start, chart re-centred)
    SpectrumParam           spectra in [lo, 1] with lambda_max = 1 and lambda_min = lo, i.e. a
                            prescribed (or bounded) Michelson contrast
    PositiveParam           X = U diag(lambda) U^*  with such a spectrum
    SelfAdjointParam        X = U diag(s * lambda) U^*  with fixed random signs s per start

    best = optimize.extremize(objective, params, maximize=True)

Initialisation (`init`, for a fraction `init_fraction` of the starts; the default 'haar' is unchanged): the extremes
of many trace functionals sit at aligned eigenbases, a set of Haar measure zero, and in a real algebra a unitary chart
U0 exp(G - G^T) never leaves the connected component (det = +-1 per sector) of U0. `structured_unitary` provides
'diagonal' (random signs on the diagonal), 'pairs' (2x2 rotations of consecutive basis vectors at random angles) and
'oriented' (Haar, with det U0_c = +1 in every sector of a real algebra) starting points; `spectrum_init` gives the
structured starts two-level spectra. With X on 'diagonal' and Y on 'pairs', the relative eigenbasis of X and Y is
a direct sum of 2x2 rotations, the structure of the maximisers of Exp. 1 and Exp. 3 (paper, Section 6), in every
dimension.

The objective maps the current operators to one value per start, shape (batch,). The search is
local: the result is a lower bound for a supremum (an upper bound for an infimum) and should be
checked with more starts / steps.

structured_unitary¶

def structured_unitary(alg: TypeIAlgebra, batch_size: int, kind: str='pairs') -> torch.Tensor

Block-diagonal unitaries (batch, C, k_max, k_max) (zero on the padding) as starting points of a search:
  'haar'      Haar-random in every sector (random_unitary_operator);
  'diagonal'  diag(+-1) with independent random signs;
  'pairs'     R diag(+-1): R rotates the consecutive basis vectors (e_0, e_1), (e_2, e_3), ... of every sector
              by independent angles uniform in [0, pi) (an odd sector keeps its last vector), with random signs;
  'rotations' R alone (det = +1 in every sector);
  'oriented'  Haar-random, and in a real algebra with det = +1 in every sector (first column negated where
              det = -1); in a complex algebra the same as 'haar'.

Param¶

class Param()

Base class: a differentiable parametrisation of a batch of operators.

Param.parameters¶

def parameters(self) -> List[torch.Tensor]

Param.recenter¶

def recenter(self)

Move the reference point to the current point (default: nothing).

UnitaryParam¶

class UnitaryParam(alg: TypeIAlgebra, batch_size: int, base: Optional[torch.Tensor]=None, init: str='haar', init_fraction=1.0)

Block-diagonal unitaries U = U0 exp(G - G^*), one per start; U0 Haar-random by default, or a structured start
(`init`, see structured_unitary) for a random fraction `init_fraction` of the starts.

UnitaryParam.parameters¶

def parameters(self)

UnitaryParam.matrix¶

def matrix(self) -> torch.Tensor

U as a tensor; computed once per parameter state (an objective that uses U several times, e.g. X and
|X| of a PositiveParam, pays one matrix exponential).

UnitaryParam.operator¶

def operator(self) -> Operator

UnitaryParam.recenter¶

def recenter(self)

SpectrumParam¶

class SpectrumParam(alg: TypeIAlgebra, batch_size: int, delta: Optional[torch.Tensor]=None, delta_range=(0.0, 0.999), pin_sector: Optional[int]=None)

Spectra (per block, padded) in [lo, 1] with the global maximum 1 and minimum lo, so that the
Michelson contrast equals delta = (1 - lo) / (1 + lo).

delta: tensor (batch,) of prescribed contrasts, or None to optimise the contrast as well
within [delta_range[0], delta_range[1]]. The maximum 1 and the minimum lo are pinned to two
eigenvalue slots; the other eigenvalues lie in [lo, 1] (through (1 + sin theta) / 2). By default
the two slots are drawn at random per start among all active slots (the extremes may lie in the
same sector or in different ones), so that a multistart search covers every placement and the
union over placements is dense in the constraint set. With `pin_sector` both extremes sit in that
sector (which needs k >= 2) for every start; for C >= 2 this is a restriction.

SpectrumParam.parameters¶

def parameters(self)

SpectrumParam.delta¶

def delta(self) -> torch.Tensor

SpectrumParam.eigenvalues¶

def eigenvalues(self) -> torch.Tensor

(batch, C, k_max), zero on padding.

PositiveParam¶

class PositiveParam(alg: TypeIAlgebra, batch_size: int, delta: Optional[torch.Tensor]=None, delta_range=(0.0, 0.999), pin_sector: Optional[int]=None, init: str='haar', init_fraction=1.0, spectrum_init: Optional[str]=None)

Positive operators X = U diag(lambda) U^* with a prescribed or bounded Michelson contrast. The eigenbasis
U = U0 exp(G - G^*) starts at a Haar-random U0, or at structured_unitary(init) for a random fraction
`init_fraction` of the starts (`self.structured` marks them).

PositiveParam.parameters¶

def parameters(self)

PositiveParam.recenter¶

def recenter(self)

PositiveParam.operator¶

def operator(self) -> Operator

PositiveParam.abs_operator¶

def abs_operator(self) -> Operator

|X| = U diag(lambda) U^*, exact (no SVD, so gradients stay finite).

PositiveParam.contrast¶

def contrast(self) -> torch.Tensor

SelfAdjointParam¶

class SelfAdjointParam(alg: TypeIAlgebra, batch_size: int, delta: Optional[torch.Tensor]=None, delta_range=(0.0, 0.999), pin_sector: Optional[int]=None, init: str='haar', init_fraction=1.0, spectrum_init: Optional[str]=None)

Self-adjoint X = U diag(s * lambda) U^* where |X| has a prescribed or bounded contrast and the
signs s = +-1 are drawn at random per start (discrete variables; the multistart covers the sign patterns, and
discrete_step flips single signs between rounds). For the starts with a structured eigenbasis (init != 'haar'),
the signs alternate within the consecutive pairs of slots (one + and one - per pair, in random order), the sign
pattern of the maximisers of Exp. 3 (paper, Section 6).

SelfAdjointParam.discrete_step¶

def discrete_step(self, evaluate: Callable[[], torch.Tensor], maximize: bool, trials: int=8)

Local search over the signs: flip the sign of one random active eigenvalue per start and keep the
flip where the objective improves (between the gradient rounds of extremize).

extremize¶

def extremize(objective: Callable[[], torch.Tensor], params: Sequence[Param], maximize: bool=True, rounds: int=10, steps: int=50, lr: float=0.05, decay: float=0.6, progress: bool=False) -> torch.Tensor

Gradient ascent (maximize=True) or descent on objective() -> tensor (batch,), every batch
element being an independent start. Adam with learning rate lr * decay**round; unitary charts
are re-centred after every round. Returns, per start, the best objective value attained at any
iterate (every iterate is an admissible point; no gradient).