STAlign alignment tutorial

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import torch
import plotly
import requests
from st_loading_functions import load_mHypothalamus, load_DLPFC
import time
import os
import scanpy as sc

from STalign import STalign

save_dir_time = './results'
iters = 5

DLPFC

"""DLPFC"""
section_ids_list = [['151507', '151508'], ['151508', '151509'], ['151509', '151510'], ['151669', '151670'], ['151670', '151671'], ['151671', '151672'], ['151673', '151674'], ['151674', '151675'], ['151675', '151676']]
run_times = []
for iter_ in range(iters):
    for section_ids in section_ids_list:
        dataset = section_ids[0] + '_' + section_ids[1]
        start_time = time.time()

        slice1 = load_DLPFC(root_dir="../benchmarking_data/DLPFC12", section_id=section_ids[0])
        slice2 = load_DLPFC(root_dir="../benchmarking_data/DLPFC12", section_id=section_ids[1])

        df1 = slice1.obs
        df2 = slice2.obs

        # print(slice1.obsm['spatial'])

        xI = np.array(slice1.obsm['spatial'].T[0])
        yI = np.array(slice1.obsm['spatial'].T[1])
        xJ = np.array(slice2.obsm['spatial'].T[0])
        yJ = np.array(slice2.obsm['spatial'].T[1])

        XI,YI,I,fig = STalign.rasterize(xI,yI,dx=3,blur=1.5)
        XJ,YJ,J,fig = STalign.rasterize(xJ,yJ,dx=3,blur=1.5)

        if torch.cuda.is_available():
            device = 'cuda:0'
        else:
            device = 'cpu'

        # keep all other parameters default
        params = {
                    'niter': 5000,
                    'device':device,
                    'epV': 50
                }


        out = STalign.LDDMM([YI,XI],I,[YJ,XJ],J,**params)

        A = out['A'].detach().cpu()
        v = out['v'].detach().cpu()
        xv = []
        for e in out['xv']:
            xv.append(e.detach().cpu())

        tpointsI= STalign.transform_points_source_to_target(xv,v,A, np.stack([yI, xI], 1))

        print("results")
        print(tpointsI)
        #switch tensor from cuda to cpu for plotting with numpy
        if tpointsI.is_cuda:
            tpointsI = tpointsI.cpu()

        #switch from row column coordinates (y,x) to (x,y)
        xI_LDDMM = tpointsI[:,1]
        yI_LDDMM = tpointsI[:,0]

        df3 = pd.DataFrame(

            {

                "aligned_x": xI_LDDMM,

                "aligned_y": yI_LDDMM,

            },


        )
        df3.index = df1.index

        results = pd.concat([df1, df3], axis=1)
        results.to_csv('./results/' + dataset + '_' + str(0) + '.csv')

        end_time = time.time()
        run_times.append(end_time - start_time)

MHypo

import numpy as np
import anndata

def find_upper_left_spot(spatial_coords):
    upper_left_index = np.argmin(np.sum(spatial_coords, axis=1))
    print(upper_left_index)
    return spatial_coords[upper_left_index]

def align_spatial_coordinates(data1, data2):
    upper_left1 = find_upper_left_spot(data1.obsm['spatial'])
    upper_left2 = find_upper_left_spot(data2.obsm['spatial'])

    shift = upper_left1 - upper_left2

    data2.obsm['spatial'] = data2.obsm['spatial'] + shift

    return data2

section_ids_list = [['-0.04', '-0.09'], ['-0.09', '-0.14'], ['-0.14', '-0.19'], ['-0.19', '-0.24']]
run_times = []
for iter_ in range(iters):
    for section_ids in section_ids_list:
        dataset = section_ids[0] + '_' + section_ids[1]
        start_time = time.time()

        slice1 = load_mHypothalamus(root_dir="../benchmarking_data/mHypothalamus", section_id=section_ids[0])
        slice2 = load_mHypothalamus(root_dir="../benchmarking_data/mHypothalamus", section_id=section_ids[1])
        slice2 = align_spatial_coordinates(slice1, slice2)
        df1 = slice1.obs
        df2 = slice2.obs

        # print(slice1.obsm['spatial'])

        xI = np.array(slice1.obsm['spatial'].T[0])
        yI = np.array(slice1.obsm['spatial'].T[1])
        xJ = np.array(slice2.obsm['spatial'].T[0])
        yJ = np.array(slice2.obsm['spatial'].T[1])

        XI,YI,I,fig = STalign.rasterize(xI,yI,dx=3,blur=1.5)
        XJ,YJ,J,fig = STalign.rasterize(xJ,yJ,dx=3,blur=1.5)

        if torch.cuda.is_available():
            device = 'cuda:0'
        else:
            device = 'cpu'

        # keep all other parameters default
        params = {
                    'niter': 5000,
                    'device':device,
                    'epV': 50
                }


        out = STalign.LDDMM([YI,XI],I,[YJ,XJ],J,**params)

        A = out['A'].detach().cpu()
        v = out['v'].detach().cpu()
        xv = []
        for e in out['xv']:
            xv.append(e.detach().cpu())

        tpointsI= STalign.transform_points_source_to_target(xv,v,A, np.stack([yI, xI], 1))

        print("results")
        print(tpointsI)
        #switch tensor from cuda to cpu for plotting with numpy
        if tpointsI.is_cuda:
            tpointsI = tpointsI.cpu()

        #switch from row column coordinates (y,x) to (x,y)
        xI_LDDMM = tpointsI[:,1]
        yI_LDDMM = tpointsI[:,0]

        df3 = pd.DataFrame(

            {

                "aligned_x": xI_LDDMM,

                "aligned_y": yI_LDDMM,

            },


        )
        df3.index = df1.index

        results = pd.concat([df1, df3], axis=1)
        results.to_csv('./results/' + dataset + '_' + str(0) + '.csv')

        end_time = time.time()
        run_times.append(end_time - start_time)