ObjSim/python/objsimpy/idl.py

# -*- coding: utf-8 -*-
''' python equivalents for some IDL functions
'''

import numpy as np

def main():
    n=50
    print(dist_vals(n))
    d = dist(n, 2*n)
    print(d)
    import pylab as plt
    plt.imshow(d)
    plt.show()

def dist(nx, ny=None, xlim=None, ylim=None):
    ''' like IDLs dist function: Calculates a 2d eucleedian distance array.        
    '''
    if ny is None:
        ny = nx
    if xlim is None:
        xvals = dist_vals(nx)
    else:
        xvals = np.linspace(xlim[0], xlim[1],nx)
        
    if ylim is None:
        yvals = dist_vals(ny)
    else:
        yvals = np.linspace(ylim[0], ylim[1],ny)
        
    x,y = np.meshgrid(xvals, yvals)
    d = np.sqrt(x**2 + y**2)
    return d

def dist_vals(n):
    maxval = n/2
    vals = np.linspace(0, maxval, maxval+1)
    if n % 2 == 0:
        return np.hstack([vals, vals[1:-1][::-1]])
    else:
        return np.hstack([vals, vals[1:][::-1]])

    

if __name__=="__main__":
    main()