scmkl.calculate_z

  1import numpy as np
  2import scipy
  3import anndata as ad
  4
  5from scmkl.tfidf_normalize import tfidf_train_test
  6from scmkl.estimate_sigma import est_group_sigma, get_batches
  7from scmkl.data_processing import process_data, get_group_mat, sample_cells
  8from scmkl.projections import gaussian_trans, laplacian_trans, cauchy_trans
  9
 10
 11def check_for_nan(adata: ad.AnnData):
 12    """
 13    Ensures only valid values are in training and test matrices.
 14
 15    Parameters
 16    ----------
 17    adata : ad.AnnData
 18        Object with `'Z_train'` and `'Z_test'` keys in `.uns` 
 19        attribute.
 20
 21    Returns
 22    -------
 23    None
 24    """
 25    n_nans = np.sum(np.isnan(adata.uns['Z_train']))
 26    n_nans += np.sum(np.isnan(adata.uns['Z_test']))
 27
 28    if n_nans:
 29        raise ValueError(
 30            "Some values in Z matrix are type `np.nan`. This is likely "
 31            "due to a small kernel width or invalid values in input Z matrix."
 32            )
 33    
 34    return None
 35
 36
 37def check_for_inf(adata: ad.AnnData):
 38    """
 39    Ensures only valid values are in training and test matrices.
 40
 41    Parameters
 42    ----------
 43    adata : ad.AnnData
 44        Object with `'Z_train'` and `'Z_test'` keys in `.uns` 
 45        attribute.
 46
 47    Returns
 48    -------
 49    None
 50    """
 51    n_infs = np.sum(np.isinf(adata.uns['Z_train']))
 52    n_infs += np.sum(np.isinf(adata.uns['Z_test']))
 53
 54    if n_infs:
 55        raise ValueError(
 56            "Some values in Z matrix are type `np.inf`. This is likely "
 57            "due to input matrix containing negative values."
 58            )
 59    
 60    return None
 61
 62
 63def get_z_indices(m, D):
 64    """
 65    Takes the number associated with the group as `m` and returns the 
 66    indices for cos and sin functions to be applied.
 67
 68    Parameters
 69    ----------
 70    m : int
 71        The chronological number of the group being processed.
 72
 73    D : int
 74        The number of dimensions per group.
 75
 76    Returns
 77    -------
 78    cos_idx, sin_idx : np.ndarray, np.ndarray
 79        The indices for cos and sin projections in overall Z matrix.
 80    """
 81    x_idx = np.arange(m*2*D ,(m + 1)*2*D)
 82    cos_idx = x_idx[:len(x_idx)//2]
 83    sin_idx = x_idx[len(x_idx)//2:]
 84
 85    return cos_idx, sin_idx
 86
 87
 88def calc_groupz(X_train, X_test, adata, D, sigma, proj_func, dtype):
 89    """
 90    Calculates the Z matrix for grouping.
 91
 92    Parameters
 93    ----------
 94    X_train : np.ndarray
 95        The filtered data matrix to calculate train Z mat for.
 96    
 97    X_test : np.ndarray
 98        The filtered data matrix to calculate test Z mat for.
 99
100    adata : anndata.AnnData 
101        AnnData object containing `seed_obj` in `.uns` attribute.
102
103    D : int
104        Number of dimensions per grouping.
105
106    sigma : float
107        Kernel width for grouping.
108
109    proj_func : function
110        The projection direction function to be applied to data.
111
112    Returns
113    -------
114    train_projections, test_projections : np.ndarray, np.ndarray
115        Training and testing Z matrices for group.
116    """  
117    if scipy.sparse.issparse(X_train):
118        X_train = X_train.toarray().astype(np.float16)
119        X_test = X_test.toarray().astype(np.float16)
120
121    W = proj_func(X_train, sigma, adata.uns['seed_obj'], D)
122
123    train_projection = np.matmul(X_train, W, dtype=dtype)
124    test_projection = np.matmul(X_test, W, dtype=dtype)
125
126    if np.isnan(train_projection).any() or np.isnan(test_projection).any():
127        train_projection = np.matmul(X_train, W, dtype=np.float64)
128        test_projection = np.matmul(X_test, W, dtype=np.float64)
129
130    return train_projection, test_projection
131
132
133def calculate_z(adata, n_features=5000, batches=10, 
134                batch_size=100) -> ad.AnnData:
135    """
136    Function to calculate Z matrices for all groups in both training 
137    and testing data.
138
139    Parameters
140    ----------
141    adata : ad.AnnData
142        created by `scmkl.create_adata()` with `adata.uns.keys()`: 
143        `'train_indices'`, and `'test_indices'`. 
144
145    n_features : int
146        Number of random feature to use when calculating Z; used for 
147        scalability.
148
149    batches : int
150        The number of batches to use for the distance calculation.
151        This will average the result of `batches` distance calculations
152        of `batch_size` randomly sampled cells. More batches will converge
153        to population distance values at the cost of scalability.
154
155    batch_size : int
156        The number of cells to include per batch for distance
157        calculations. Higher batch size will converge to population
158        distance values at the cost of scalability.
159        If `batches*batch_size > num_training_cells`,
160        `batch_size` will be reduced to 
161        `int(num_training_cells / batches)`.
162
163    Returns
164    -------
165    adata : ad.AnnData
166        `adata` with Z matrices accessible with `adata.uns['Z_train']` 
167        and `adata.uns['Z_test']`.
168
169    Examples
170    --------
171    >>> adata = scmkl.estimate_sigma(adata)
172    >>> adata = scmkl.calculate_z(adata)
173    >>> adata.uns.keys()
174    dict_keys(['Z_train', 'Z_test', 'sigmas', 'train_indices', 
175    'test_indices'])
176    """
177    # Number of groupings taking from group_dict
178    n_pathway = len(adata.uns['group_dict'].keys())
179    D = adata.uns['D']
180
181    sq_i_d = np.sqrt(1/D)
182
183    # Capturing training and testing sizes
184    train_len = len(adata.uns['train_indices'])
185    test_len = len(adata.uns['test_indices'])
186
187    if batch_size * batches > train_len:
188        old_batch_size = batch_size
189        batch_size = int(train_len/batches)
190        print("Specified batch size required too many cells for "
191                "independent batches. Reduced batch size from "
192                f"{old_batch_size} to {batch_size}")
193
194    if 'sigma' not in adata.uns.keys():
195        n_samples = np.min((2000, train_len))
196        sample_range = np.arange(n_samples)
197        batch_idx = get_batches(sample_range, adata.uns['seed_obj'], 
198                                batches=batches, batch_size=batch_size)
199        sigma_indices = sample_cells(adata.uns['train_indices'], n_samples, adata.uns['seed_obj'])
200
201    # Create Arrays to store concatenated group Zs
202    # Each group of features will have a corresponding entry in each array
203    n_cols = 2*adata.uns['D']*n_pathway
204    Z_train = np.zeros((train_len, n_cols), dtype=np.float16)
205    Z_test = np.zeros((test_len, n_cols), dtype=np.float16)
206
207
208    # Setting kernel function 
209    match adata.uns['kernel_type'].lower():
210        case 'gaussian':
211            proj_func = gaussian_trans
212        case 'laplacian':
213            proj_func = laplacian_trans
214        case 'cauchy':
215            proj_func = cauchy_trans
216
217
218    # Loop over each of the groups and creating Z for each
219    z_dtype = np.float16
220    sigma_list = list()
221    for m, group_features in enumerate(adata.uns['group_dict'].values()):
222
223        n_group_features = len(group_features)
224
225        X_train, X_test = get_group_mat(adata, n_features, group_features, 
226                                        n_group_features, process_test=True)
227        
228        if adata.uns['tfidf']:
229            X_train, X_test = tfidf_train_test(X_train, X_test)
230
231        # Data filtering, and transformation according to given data_type
232        # Will remove low variance (< 1e5) features regardless of data_type
233        # If scale/transform data depending on .uns values
234        X_train, X_test = process_data(X_train=X_train, X_test=X_test, 
235                                       scale_data=adata.uns['scale_data'], 
236                                       transform_data=adata.uns['transform_data'],
237                                       return_dense=True, 
238                                       add_ones=adata.uns['add_ones'])    
239
240        # Getting sigma
241        if 'sigma' in adata.uns.keys():
242            sigma = adata.uns['sigma'][m]
243        else:
244            sigma = est_group_sigma(adata, X_train, n_group_features, 
245                                    n_features, batch_idx=batch_idx)
246            sigma_list.append(sigma)
247            
248        assert sigma > 0, "Sigma must be more than 0"
249        train_projection, test_projection = calc_groupz(X_train, X_test, 
250                                                        adata, D, sigma, 
251                                                        proj_func, z_dtype)
252
253        if train_projection.dtype != z_dtype:
254            print("Converting Z dtype from np.float16 to np.float64", 
255                  flush=True)
256            z_dtype = train_projection.dtype
257            Z_train = Z_train.astype(z_dtype)
258            Z_test = Z_test.astype(z_dtype)
259
260        # Store group Z in whole-Z object
261        # Preserves order to be able to extract meaningful groups
262        cos_idx, sin_idx = get_z_indices(m, D)
263
264        Z_train[0:, cos_idx] = np.cos(train_projection, dtype=z_dtype)
265        Z_train[0:, sin_idx] = np.sin(train_projection, dtype=z_dtype)
266
267        Z_test[0:, cos_idx] = np.cos(test_projection, dtype=z_dtype)
268        Z_test[0:, sin_idx] = np.sin(test_projection, dtype=z_dtype)
269        
270    adata.uns['Z_train'] = Z_train*sq_i_d
271    adata.uns['Z_test'] = Z_test*sq_i_d
272
273    if 'sigma' not in adata.uns.keys():
274        adata.uns['sigma'] = np.array(sigma_list)
275
276    check_for_nan(adata)
277    check_for_inf(adata)
278
279    return adata