How to sort 3d numpy masked array on time axis at each grid i.e at specific latitude and longitude for climate dataset

Question

I have a NetCDF file of Climate dataset having 3D structure with a shape of 20 * 445 * 445 as (time, latitude, longitude)

I have read it as numpy.ma.core.MaskedArray using netCDF4 library

I want to sort this NumPy 3D-array data on the time axis at each grid i.e.

at any specific grid (latitude,longitude) the value should be in either ascending or descending order

I have tried the numpy.argsort function but the range of values at any grid/pixel exceeds the original value of dataset, so its not working i.e. the original dataset has values between 7 to 16, but after sorting, the values in dataset ranges from 1 to 19.

the code is available at the Google Collab, and the NetCDF is available here

import netCDF4 as nc
from netCDF4 import Dataset,num2date,date2num
import sys,os,numpy as np,logging
from datetime import date, datetime, timedelta
import matplotlib.pyplot as plt

fileName = 'some_file'
data = Dataset(fileName, 'r')

temp_data = data.variables['Band1'][:,:,:]

Plots of original Dataset

for i in range(0,20):
    data_to_plot = temp_data[i]
    
    fig, ax = plt.subplots(figsize = (4,4))
    ax.set_title('Layer '+ str(i))
    cax = ax.imshow(data_to_plot, cmap = plt.cm.Accent)
    
    cbar = plt.colorbar(cax,
                        orientation='vertical',
                        fraction=0.045, 
                        pad=0.05)
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')

sorted_temp_data = np.argsort(temp_data, axis=0)

Plots of Sorted Datasets

for i in range(0,20):
    data_to_plot = sorted_temp_data[i]
    
    fig, ax = plt.subplots(figsize = (4,4))
    ax.set_title('Layer '+ str(i))
    cax = ax.imshow(data_to_plot, cmap = plt.cm.Accent)
    
    cbar = plt.colorbar(cax,
                        orientation='vertical',
                        fraction=0.045, 
                        pad=0.05)
    
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')

How to do so?

Answer 1

Numpy's argsort function doesnt return sorted array. It returns the index of sorted array. You can directly sort the array along the time axis (assuming time axis is the 0th axis) as follows:

np.sort(temp_data, axis=0)

If you want to use argsort the code should be like this:

idx = np.argsort(temp_data, axis=0)
sort_data = np.take_along_axis(temp_data, idx,axis=0)