spook package

spook module

class spook.base.SpookBase(B, A, mode='raw', G=None, lsparse=None, lsmooth=None, Bsmoother='laplacian', Asmoother='laplacian', normalize=True)

Bases: object

Base class for Spook solvers

property AGscale

Product s_A * s_G, where s_A and s_G are the normalization factors of A and G.

property AGtAG

Tensor product of AtA and GtG Some children classes need this tensor itself, some need its upper triangular part only, so I make it a non-cache property A child class can cache it/its upper triangular part if necessary

Asm()

Build the smoothness operator for the dimension of A

property Na

Dimension of single-shot A

property NaTuple

Shape of single-shot A. Most frequently, it is (Nw,) But in the case of delay binning, it is (Nw, Nt)

property Ng

(Dimension of X) / Na

accumulate(AtA_batch, Bcontracted_batch)

getXopt(lsparse=None, lsmooth=None, orig_scale=True)

Solve for the optimal X, :param lsparse: The regularization hyperparameters :type lsparse: float|None, lsmooth: float|tuple|None :param orig_scale: whether to scale it back to the original scale :type orig_scale: bool

normalizeAG(normalize)

residueL2(X=None)

Calculate the rmse of the linear model given variable X, in the original (unscaled) problem rmse = |(A otimes G)X - B|_2 :param X: Optimization (unscaled) variable. If None, use the last solved X :type X: np.ndarray

save_prectr(inplace=False)

Export the pre-contracted arrays into a dictionary

set_btb(btb: float)

Set the constant term in the quadratic form For 1D problem, btb is the sum of b^2 across all shots For 2D problem, btb is the sum of b^2 across all shots and all KE bins

smoothness(X=None, dim='w')

Evaluate the smoothness of variable X, along the dimension dim. The smoothness is defined as sqrt(X.T @ LL @ X), where LL is the smoothness operator

by default LL is the square of finite difference Laplacian operator. No lsmooth* is applied.

Parameters:

• X (the solution to evaluate)

• dim ('w' or 'b' or 't') – ‘w’: along the photon frequency axis ‘b’: along the interested property axis ‘t’: along the delay axis (if applicable)

smoothness_drop_boundaries = True

solve(lsparse=None, lsmooth=None)

ONLY for the base class For every derived class, PLEASE REDEFINE!

sparsity(X=None)

L2 norm of the solution. No square, no lsparse*.

update_lsmooth(lsmooth)

update_lsparse(lsparse)

To be redefined in each derived class

verbose = False

class spook.quad_program.SpookL1(B, A, mode='raw', G=None, lsparse=1, lsmooth=(0.1, 0.1), osqp_settings={}, **kwargs)

Bases: SpookQPBase

L1 sparsity w/o non-negativity. This is lasso, which requires auxilary variables

calcPtriu()

setupProb(col=None)

Create a new OSQP problem Upper triangular part of self._P is used in OSQP.setup(), regardless of whether self._P is dense or sparse

solve(lsparse=None, lsmooth=None)

Solve the quadratic program

sparsity(X=None)

L1 norm of X. No lsparse* applied.

update_lsparse(lsparse)

To be redefined in each derived class

class spook.quad_program.SpookPos(B, A, mode='raw', G=None, lsparse=1, lsmooth=(0.1, 0.1), osqp_settings={}, **kwargs)

Bases: SpookQPBase

Nonnegativity constraint Definitely a Quadratic Program

initProb()

Create a new OSQP problem Upper triangular part of self._P is used in OSQP.setup(), regardless of whether self._P is dense or sparse Child class call this to get bounds and OSQP instance

class spook.quad_program.SpookPosL1(B, A, mode='raw', G=None, lsparse=1, lsmooth=(0.1, 0.1), osqp_settings={}, **kwargs)

Bases: SpookPos

Positivity + L1 sparsity L1 sparsity is just a linear term

calcPtriu()

setupProb(col=None)

Create a new OSQP problem Upper triangular part of self._P is used in OSQP.setup(), regardless of whether self._P is dense or sparse

solve(lsparse=None, lsmooth=None)

Solve the quadratic program

sparsity(X=None)

L1 norm of X. No lsparse* applied.

update_lsparse(lsparse)

To be redefined in each derived class

class spook.quad_program.SpookPosL2(B, A, mode='raw', G=None, lsparse=1, lsmooth=(0.1, 0.1), osqp_settings={}, **kwargs)

Bases: SpookPos

Non-negativity + L2^2 sparsity, i.e. Ridge w/ Non-negativity constraints L1 sparsity is just a linear term

calcPtriu()

setupProb(col=None)

Create a new OSQP problem Upper triangular part of self._P is used in OSQP.setup(), regardless of whether self._P is dense or sparse

solve(lsparse=None, lsmooth=None)

Solve the quadratic program

update_lsparse(lsparse)

To be redefined in each derived class

class spook.quad_program.SpookQPBase(B, A, mode='raw', G=None, lsparse=1, lsmooth=(0.1, 0.1), osqp_settings={}, **kwargs)

Bases: SpookBase

calc_total_smoother_triu()

CALCULATE Total Smoothness regularization operator lsma * (La.T @ La otimes Ib) + lsmb * (Ia otimes Bsmoother)

This method always recalculates

Returns:

The upper triangular part of tensor product matrix, so DON’T use if the solver is vectorized over dimension b

qhalf(col=None)

set_polish(polish_bool=True)

solve(lsparse=None, lsmooth=None)

Solve the quadratic program

update_lsmooth(lsmooth)

verbose = False

class spook.lin_solve.SpookLinSolve(B, A, mode='raw', G=None, lsparse=1, lsmooth=(0.1, 0.1), **kwargs)

Bases: SpookBase

Spooktroscopy that involves only linear eq solving This means: no positivity constraint L2 squared sparsity

setupProb()

solve(lsparse=None, lsmooth=None)

Solve the linear system, since there is only a Tikhonov regularization but no constraint.

update_lsmooth(lsmooth)

update_lsparse(lsparse)

To be redefined in each derived class

verbose = False

spook.utils module

spook.utils.calcL2fromContracted(Xo, AtA, Bcontracted, trBtB, GtG=None)

Calculates the residual L2 norm from the contracted tensors. :param Xo: :type Xo: ndarray, value of the optimization variable :param AtA: :type AtA: ndarray, AtA matrix :param Bcontracted: :type Bcontracted: ndarray, AtB matrix if G is None, otherwise Bcontracted = AtB @ G :param trBtB: :type trBtB: float, constant in the mean square error :param GtG: :type GtG: ndarray, (optional) GtG matrix. If None, it means G=identity

spook.utils.count_delaybin(at_iter)

spook.utils.dict_allsqsum(dictB)

spook.utils.dict_innerprod(dictA, dictB, Aroi=None)

Inner product of two tensors represented as dictionaries, with the contracted dimension being the keys.

spook.utils.eval_Ng(b_iter)

spook.utils.eval_Nw(at_iter)

spook.utils.iso_struct(csc_mata, csc_matb)

Determine whether two csc sparse matrices share the same structure

spook.utils.laplacian1D_S(N)

spook.utils.laplacian_square_S(N, drop_bound)

spook.utils.matricize_tensor_bykey(dct, ky_list, roi=None)

spook.utils.poisson_nll(pred, data)

spook.utils.soft_poisson_nll(pred, data, p=0.01)

spook.utils.worth_sparsify(arr)

spook.xval module

class spook.xval.XValidation(solver_class, datasets, **kwargs)

Bases: object

Class for cross-validation of a solver.

calc_residual(hyperparams, dset='val', avg=True)

Validate the hyperparameters on the chosen set. :param hyperparams: The hyperparameters to validate.

Each tuple must be in the form of (lsparse, lsmooth) 2D array will be interpreted as a list of (lsparse, *lsmooth)

Parameters:

• dset (str) – which dataset to validate on. Must be either ‘train’ or ‘val’

• avg (bool) – If True, return the average of the residuals. Otherwise, return the residuals for each train/val split

Returns:

The L2 norm (no square) of the residuals. If avg is True, return the root-mean-squared of the L2 norm over the k splits

Return type:

float | np.ndarray

fit(lsparse=None, lsmooth=None)

property k

spook.xval.parse_atb(atb)

spook.xval.parse_btb(btb)

spook.vmi_special module

spook.contraction_utils module

Note

Function adaptive_contraction in this module serves in the raw mode of the spook solver. It is not parallelized, so it is slow for large datasets. For TMO users, a good alternative is the outer-product accumulation application in the tmo-preproc repo.

spook.contraction_utils.adaptive_contraction(A, B, a_slice=None, b_slice=None, keep_dims=False)

Contract two tensors A, B along axis0, generalizing numpy.tensordot in terms of memory limit. :param A: two tensors stores as arrays/lists/dictionaries.

Arrays and lists can be mixed-and-matched, but if one of A,B is a dictionary, the other has to be too, with the same set of keys.

Parameters:

• B – two tensors stores as arrays/lists/dictionaries. Arrays and lists can be mixed-and-matched, but if one of A,B is a dictionary, the other has to be too, with the same set of keys.

• a_slice – optional slicing objects that can be applied to each singleshot data in A, B, respectively

• b_slice – optional slicing objects that can be applied to each singleshot data in A, B, respectively

• keep_dims – controlling flag regarding whether the dimensions of the final tensor is kept. If False, the final tensor will be reshaped into a matrice as if A, B were two matrices with axis0 being the shot id while axis1 being the rest.

spook.contraction_utils.adaptive_dot(B, G)

spook.contraction_utils.allsqsum(A, a_slice=None)

spook.contraction_utils.check_matchness(A, B)

spook.contraction_utils.check_sliceobj(obj)

spook.contraction_utils.shots_asarray(shots_iterable, shots_select=None, singleshot_slice=None)

Iterate through an iterable object shots_iterable, apply optional singleshot_slice to each singleshot array, then stack up the sliced arrays into an array, whose axis0 is shot id.

Use shots_select to specify which shots to stack into the matrix:

If the iterable is a dictionary, shots_select is a list of keys If the iterable is a list of np.ndarray, shots_select is a list of indices or a slice

This is a lower-level function that only accepts formated singleshot_slice object, if not None