scmkl.optimize_sparsity

  1import numpy as np
  2import anndata as ad
  3
  4from scmkl.train_model import train_model
  5from scmkl.test import find_selected_groups
  6
  7
  8def optimize_sparsity(adata: ad.AnnData, group_size: int | None=None, starting_alpha = 1.9, 
  9                      increment = 0.2, target = 1, n_iter = 10):
 10    """
 11    Iteratively train a grouplasso model and update alpha to find the 
 12    parameter yielding the desired sparsity.
 13    
 14    Parameters
 15    ----------
 16    adata : ad.AnnData
 17        `ad.AnnData` with `'Z_train'` and `'Z_test'` in 
 18        `adata.uns.keys()`.
 19
 20    group_size : None | int
 21        Argument describing how the features are grouped. If `None`, 
 22        `2 * adata.uns['D']` will be used. For more information see 
 23        [celer documentation](https://mathurinm.github.io/celer/
 24        generated/celer.GroupLasso.html).
 25
 26    starting_alpha : float
 27        The alpha value to start the search at.
 28    
 29    increment : float 
 30        Amount to adjust alpha by between iterations.
 31    
 32    target : int
 33        The desired number of groups selected by the model.
 34
 35    n_iter : int
 36        The maximum number of iterations to run.
 37            
 38    Returns
 39    -------
 40    sparsity_dict : dict
 41        Tested alpha as keys and the number of selected groups as 
 42        the values.
 43        
 44    alpha : float
 45        The alpha value yielding the number of selected groups closest 
 46        to the target.
 47
 48    Examples
 49    --------
 50    >>> sparcity_dict, alpha = scmkl.optimize_sparsity(adata, 
 51    ...                                                target = 1)
 52    >>>
 53    >>> alpha
 54    0.01
 55
 56    See Also
 57    --------
 58    celer.GroupLasso : https://mathurinm.github.io/celer/
 59    """
 60    assert increment > 0 and increment < starting_alpha, ("Choose a positive "
 61                                                          "increment less "
 62                                                          "than alpha")
 63    assert target > 0 and isinstance(target, int), ("Choose an integer "
 64                                                    "target number of groups "
 65                                                     "that is greater than 0")
 66    assert n_iter > 0 and isinstance(n_iter, int), ("Choose an integer "
 67                                                    "number of iterations "
 68                                                    "that is greater than 0")
 69
 70    if group_size == None:
 71        group_size = adata.uns['D']*2
 72
 73    sparsity_dict = {}
 74    alpha = starting_alpha
 75
 76    for _ in np.arange(n_iter):
 77        adata = train_model(adata, group_size, alpha)
 78        num_selected = len(find_selected_groups(adata))
 79
 80        sparsity_dict[np.round(alpha, 4)] = num_selected
 81
 82        if num_selected < target:
 83            #Decreasing alpha will increase the number of selected pathways
 84            if alpha - increment in sparsity_dict.keys():
 85                # Make increment smaller so the model can't go back and forth 
 86                # between alpha values
 87                increment/=2
 88            # Ensures that alpha will never be negative
 89            alpha = np.max([alpha - increment, 1e-3]) 
 90
 91        elif num_selected > target:
 92            if alpha + increment in sparsity_dict.keys():
 93                increment/=2
 94
 95            alpha += increment
 96        elif num_selected == target:
 97            break
 98
 99    # Find the alpha that minimizes the difference between target and observed
100    # number of selected groups
101    spar_idx = np.argmin([np.abs(selected - target) 
102                          for selected in sparsity_dict.values()])
103    optimal_alpha = list(sparsity_dict.keys())[spar_idx]
104    
105    return sparsity_dict, optimal_alpha