ModelVF.py

"""
 ModelVF.py

 For a general description of the whole package and how to import,
 see the head of the file PyCigale.py
 
 Usage:
 PyC.model('disk','output.fits',256,100,10,60,70)

 Models implemented so far:
 * linear: gives a velocity field with a simple linear gradient
 * shell: velocity field of a shell that is beeing looked at
   ( it turned out to be linear as well :-))
 * disk: a thin disk in DM-potential

 See the function descriptions as well.

 This file needs better integration into the package PyCigale!!
"""

from PyGalKin import *

def interact_model_linear(r, pars):
	# calculate velocities from gradient and radius
	arr = r[:].copy()
	velgrad = pars['gradient']
	sysvel = pars['system']
	arr = arr * velgrad + sysvel
	return arr


def galmod_fitlin(x, y, p):
	# returns linear vel field from coordinates x,y and fitting parameters p
	pars = {}
	pars['pa'] = p[0]
	pars['gradient'] = p[1]
	pars['system'] = p[2]
	pars['centr_offset_x'] = 0
	pars['centr_offset_y'] = 0
	
	radii = radii_from_position((x,y),pars)
	vels = interact_model_linear(radii, pars)
	return vels
	

def linfit_func(p, fjac, x=None, y=None, z=None, err=None, returnmodel=False):
	# linear galaxy rotation modelling, conform to mpfit
	model = galmod_fitlin(x, y, p)
	if returnmodel: return model
	else: return([0, z-model])
		
		
def expfit_func(p, fjac, x=None, y=None, z=None, err=None, returnmodel=False):
	# exponential galaxy flux modelling, conform to mpfit
	model = p[0]+(p[1]*N.sqrt(x**2+y**2))
	if returnmodel: return model
	else: return([0, z-model])
		
		
def powerlawfit_func2(p, fjac, x=None, y=None, z=None, err=None, returnmodel=False):		
	# exponential galaxy flux modelling, conform to mpfit
	model = p[0]+p[1]*N.log(N.sqrt(x**2+y**2))
	if returnmodel: return model
	else: return([0, z-model])
		
def powerlawfit_func(p, fjac, x=None, y=None, z=None, err=None, returnmodel=False):		
	# exponential galaxy flux modelling, conform to mpfit
	model = N.where(p[0]+p[1]*N.log(N.sqrt(x**2+y**2))<p[2],p[0]+p[1]*N.log(N.sqrt(x**2+y**2)),p[2])
	if returnmodel: return model
	else: return([0, z-model])
	
		
def gauss_fit_func(p, fjac, x=None, y=None, err=None, returnmodel=False):
	# gaussian fitting function, conform to mpfit
	model = p[0]*N.exp(-(x-p[2])**2/(2*p[1]**2))+p[3]
	if returnmodel: return model
	else: return([0, y-model])		
		

def lin_fit_func(p, fjac, x=None, y=None, err=None, returnmodel=False):
	# linear fitting function, conform to mpfit
	model = p[0]+x*p[1]
	if returnmodel: return model
	else: return([0, y-model])		
		
		


def rotation_curve(rings, outfile=None):
  """Generates a position and velocity vector from a list of rings generated
      by tilted_ring_model(). It can also write the output to a .rc-file.
      
      Usage: pos,vel = rotation_curve(rings, outfile)
      
      rings: A list of rings generated by tilted_ring_model()
      outfile: Optional filename to write the data to outfile.rc (default is None)
      pos,vel: Vectors containing position and velocity along a rotation curve.
      """
  # Create empty lists
  vel=[]
  pos=[]
  err=[]
  
  # For each ring, read the radius (scaled to pc's) and velocity
  for i in range(len(rings)):
    pos.append(float(rings[i]['radius']*rings[i]['scale']))
    vel.appens(float(rings[i]['vel']))
    err.append(float(0.0))
  
  # Write the outfile if a name is given
  if (outfile != None):
    write_rc(outfile, pos, vel, err)
  
  return pos,vel


def tilted_ring_model(data, width):
  """Creates a tilted ring-model from a velocity field.
      
      Usage: system,rings = tilted_ring_model(data, width)
      
      data: The velocity field to model
      width: The width of each ring, this will set the number of rings
      system: The system velocity
      rings: A list of dictionaries (one for each ring) with parameters on the 
      form:
        [{'centr_x', 'centr_y', 'vel', 'inclination', 'pa', 'radius', 'width', 
        'dim', 'scale'},
        {'centr_x', 'centr_y', 'vel', 'inclination', 'pa', 'radius', 'width', 
        'dim', 'scale'}, ... ]
        
        centr_x,centr_y: The center of the ring
        vel: The velocity of the ring
        inclination: The inclination of the ring
        pa: The position angle of the ring
        radius: The inner radius of the ring
        width: The width of the ring
        dim: The dimension of the array for the ring (always a square array)
        scale: The parsec/pixel-scale
      """
  # We do not want to edit a link to the data
  temp = data.copy()
  
  # Understandable variables
  dim = temp.nx()
  centr_x = temp.p['dyncen'][0]
  centr_y = temp.p['dyncen'][1]
  vel = 50
  inclination = 45
  pa = temp.p['pa']
  minimum = temp.min()
  
  # Fit the system velocity, only one parameter
  parinfo = []
  parinfo.append({'value':0.0, 'fixed':0, 'limited':[0,0],'limits':[0.0, 0.0], 'step':0.0})
  parinfo[0]['value'] = temp.mean()
  
  functkw = {'x':temp, 'y':dim, 'err':minimum}
  temp1 = mpfit.mpfit(model_system_func, parinfo=parinfo, functkw=functkw, quiet=1)
  v_system = temp1.params[0]
  
  print(v_system)
  
  # Empty list for the rings
  pars = []
  
  # Compute the rings
  for i in range(int(0.5*dim/width)):
    radius = i*width
    print('Ring: '+str(i+1)+', radius: '+str(radius))
    
    # Parameters to fit
    parinfo = []
    for j in range(5):
      parinfo.append({'value':0.0, 'fixed':0, 'limited':[0,0],'limits':[0.0, 0.0], 'step':0.0})
    
    parinfo[0]['value'] = pa
    parinfo[0]['limited'][0] = 0
    parinfo[0]['limits'][0] = 0
    parinfo[0]['limited'][1] = 0
    parinfo[0]['limits'][1] = 0
    parinfo[0]['fixed'] = 0
    
    parinfo[1]['value'] = vel
    parinfo[1]['limited'][0] = 0
    parinfo[1]['limits'][0] = 0
    parinfo[1]['limited'][1] = 0
    parinfo[1]['limits'][1] = 0
    parinfo[1]['fixed'] = 0
    
    parinfo[2]['value'] = inclination
    parinfo[2]['limited'][0] = 0
    parinfo[2]['limits'][0] = 0
    parinfo[2]['limited'][1] = 0
    parinfo[2]['limits'][1] = 0
    parinfo[2]['fixed'] = 0
    
    parinfo[3]['value'] = centr_x
    parinfo[3]['limited'][0] = 0
    parinfo[3]['limits'][0] = 0
    parinfo[3]['limited'][1] = 0
    parinfo[3]['limits'][1] = 0
    parinfo[3]['fixed'] = 0
    
    parinfo[4]['value'] = centr_y
    parinfo[4]['limited'][0] = 0
    parinfo[4]['limits'][0] = 0
    parinfo[4]['limited'][1] = 0
    parinfo[4]['limits'][1] = 0
    parinfo[4]['fixed'] = 0
    
    functkw = {'x':temp, 'y':[radius, width, dim, v_system, minimum], 'err':0}
    
    temp2 = mpfit.mpfit(model_tilted_ring_func, parinfo=parinfo, functkw=functkw, quiet=1)
    p = temp2.params
    
    # Add the ring to pars
    pars.append({'centr_x':p[3], 'centr_y':p[4], 'vel':p[1], 'inclination':p[2], 'pa':p[0], 'radius':radius, 'width':width, 'dim':dim, 'scale':temp.scale()})
    
    print(str(pars[i]))
  
  return [v_system, pars]

def model_system_func(p, fjac=None, x=None, y=None, err=None): 
  """Used by tilted_ring_model() to fit a system velocity to a velocity field.
      It will only consider points containing data, not empty points.
      
      Note: This is not supposed to be used from the command line.
      """
  # Set meaningful variable names
  data = x
  dim = y
  minimum = err
  
  # Array to store the result
  result = N.zeros((dim,dim), dtype='Float64')
  # Uniform velocity field
  vf_system = N.zeros((dim,dim), dtype='Float64') + p[0]
  # Points where there is data
  points = N.where(data > minimum)
  # Difference between model and data
  result[points] = vf_system[points] - data[points]
  result.setshape(dim*dim)
  
  status=0
  return [status, result]
  
def model_tilted_ring_func(p, fjac=None, x=None, y=None, err=None):
  """Used by tilted_ring_model() to fit a tilted ring to a velocity field.
      It will only consider points containing data (both in the velocity field 
      and the ring), not empty points.
      
      Note: This is not supposed to be used from the command line.
      """
  # Setup variables with understandable names
  data = x
  radius = y[0]
  width = y[1]
  dim = y[2]
  v_system = y[3]
  minimum = y[4]
  
  # Create a system velocity field from the known system velocity
  vf_system = N.zeros((dim,dim), dtype='Float64') + v_system

  # Create a ring and add the system velocity
  ring = create_ring(p, radius, width, dim)
  
  # Only consider points where the ring has some data and
  # where the 'data'  has data
  points1 = N.where(ring != 0)
  points2 = N.where(data > minimum)
  
  ring[points1] = ring[points1] + vf_system[points1]
  
  # The difference between the data and the model
  result_temp = N.zeros((dim,dim), dtype='Float64')
  result_temp[points1] = ring[points1] - data[points1]
  
  result = N.zeros((dim,dim), dtype='Float64')
  result[points2] = result_temp[points2]
  result.setshape(dim*dim)
  
  print(str(p))
  
  status=0
  return [status, result]

def create_ring(p, radius, width, dim):
  """Creates a tilted ring as viewed from the telescope.
      
      Usage: ring_tr = create_ring(p, radius, width, dim)
      
      p: A list of parameters for the ring [pa,vel,inclination,centr_x,centr_y]
        centr_x, centr_y: The centre of the ring
        vel: The velocity
        inclination: The inclination of the ring
        pa: The position angle of the ring
      radius: The inner radius of the ring
      width: The width of the ring
      dim: The dimension of the array (the array is always a square)
      ring_tr: A 2D-array with the tilted ring
      """
  # Set understandable variable names
  centr_x = p[3]
  centr_y = p[4]
  vel = p[1]
  inclination = p[2]
  pa = p[0]
  
  # Convert the pa to rads and transform it to the correct definition
  pa = (-pa-90)*M.pi/180
  # Convert the inclination to rads
  inclination = inclination*M.pi/180

  # Arrays needed for the computation
  r = N.zeros((dim,dim), dtype='Float64')
  phi = N.zeros((dim,dim), dtype='Float64')
  ring_non_tr = N.zeros((dim,dim), dtype='Float64')
  ring_tr = N.zeros((dim,dim), dtype='Float64')
  
  # Create an array with the angle and one with the radius in each point
  # calculated from the ring-centre.
  for x in range(dim):
    for y in range(dim):
      temp_x = x-centr_x
      temp_y = y-centr_y
      
      phi[x,y] = N.arctan2(temp_y, temp_x)
      r[x,y] = N.sqrt(temp_x**2+temp_y**2)

  # Mark points within the inner ring as t2 and points within the outer ring as t1
  t1 = N.where(r < (radius+width))
  t2 = N.where(r < radius)
  
  # Create a non transformed ring
  ring_non_tr[t1] = vel*N.cos(phi[t1])*N.sin(inclination)
  ring_non_tr[t2] = 0.0
  
  # Transform the ring to the telescope-coordinates (from the galaxy coords.)
  for x in range(dim):
    for y in range(dim):
      temp_x = x-centr_x
      temp_y = y-centr_y
      
      x_corr = temp_x*N.cos(pa) - temp_y*N.sin(pa)
      y_corr = temp_x*N.sin(pa) + temp_y*N.cos(pa)
      y_corr = y_corr/N.cos(inclination)
      
      x_corr2 = int(x_corr+centr_x)
      y_corr2 = int(y_corr+centr_y)
      
      if (x_corr2 >= 0 and x_corr2 < dim) and (y_corr2 >= 0 and y_corr2 < dim):
        ring_tr[x,y] = ring_non_tr[x_corr2, y_corr2]
  
  return ring_tr


def fit_parameters(data, model_list, box_size=3):
  """Fits a model defined in model_list to the data. You interactivly choose
      boxes of size 'box_size' to use for comparing the model and the data.
      
      An average value will be calculated for each box you choose and
      compared with the average from the same box in the model. This
      difference is minimized.
      
      The output is a model_list with the best fit.
      
      Usage: model_list_out = fit_parameters(data, model_list_in, box_size)
      
      data: The velocity field
      model_list_in: A list of models to fit and the initial values
      box_size: The size of the boxes to average (a kind of low pass filter)
      model_list_out: A list of the models and the best fit parameters
      """
  # Ask the user for boxes
  boxes = get_boxes(data)
  # Average the boxes
  boxes_avg = get_boxes_avg(data, boxes, box_size)
  
  # Fixes some obscure bug, linking was used instead of copying
  model_list = N.copy.deepcopy(model_list)
  
  # Set limits on the parameters
  parinfo = []

  for i in range( len(model_list)*10 + 1 ):
    parinfo.append({'value':0.0, 'fixed':1, 'limited':[0,0],'limits':[0.0, 0.0], 'step':0.0})
  
  parinfo[0]['value'] = model_list[0][1]['dim']
  
  for i in range(len(model_list)):
    parinfo[i*10+1]['value'] = model_list[i][1]['pa']
    parinfo[i*10+1]['limited'][0] = 1
    parinfo[i*10+1]['limits'][0] = 0
    parinfo[i*10+1]['limited'][1] = 1
    parinfo[i*10+1]['limits'][1] = 360
    parinfo[i*10+1]['fixed'] = 0
    
    parinfo[i*10+2]['value'] = model_list[i][1]['inclination']
    parinfo[i*10+2]['limited'][0] = 1
    parinfo[i*10+2]['limits'][0] = 0
    parinfo[i*10+2]['limited'][1] = 1
    parinfo[i*10+2]['limits'][1] = 90
    parinfo[i*10+2]['fixed'] = 0
    
    parinfo[i*10+3]['value'] = model_list[i][1]['exp_max']
    parinfo[i*10+3]['limited'][0] = 1
    parinfo[i*10+3]['limits'][0] = 0
    parinfo[i*10+3]['limited'][1] = 1
    parinfo[i*10+3]['limits'][1] = model_list[0][1]['dim']
    parinfo[i*10+3]['fixed'] = 1
    
    parinfo[i*10+4]['value'] = model_list[i][1]['r_max']
    parinfo[i*10+4]['limited'][0] = 1
    parinfo[i*10+4]['limits'][0] = 0
    parinfo[i*10+4]['limited'][1] = 1
    parinfo[i*10+4]['limits'][1] = model_list[0][1]['dim']
    parinfo[i*10+4]['fixed'] = 1
    
    parinfo[i*10+5]['value'] = model_list[i][1]['v_max']
    parinfo[i*10+5]['limited'][0] = 1
    parinfo[i*10+5]['limits'][0] = 0
    parinfo[i*10+5]['limited'][1] = 1
    parinfo[i*10+5]['limits'][1] = 1000
    parinfo[i*10+5]['fixed'] = 0
    
    parinfo[i*10+6]['value'] = model_list[i][1]['v_system']
    parinfo[i*10+6]['limited'][0] = 1
    parinfo[i*10+6]['limits'][0] = 0
    parinfo[i*10+6]['limited'][1] = 1
    parinfo[i*10+6]['limits'][1] = 2000
    parinfo[i*10+6]['fixed'] = 0
    
    parinfo[i*10+7]['value'] = model_list[i][1]['v_expansion']
    parinfo[i*10+7]['limited'][0] = 1
    parinfo[i*10+7]['limits'][0] = 0
    parinfo[i*10+7]['limited'][1] = 1
    parinfo[i*10+7]['limits'][1] = 1000
    parinfo[i*10+7]['fixed'] = 1
    
    parinfo[i*10+8]['value'] = model_list[i][1]['a_scale']
    parinfo[i*10+8]['limited'][0] = 1
    parinfo[i*10+8]['limits'][0] = 0
    parinfo[i*10+8]['limited'][1] = 1
    parinfo[i*10+8]['limits'][1] = 100
    parinfo[i*10+8]['fixed'] = 0
    
    parinfo[i*10+9]['value'] = model_list[i][1]['centr_offset_x']
    parinfo[i*10+9]['limited'][0] = 1
    parinfo[i*10+9]['limits'][0] = -int(0.5*model_list[0][1]['dim'])
    parinfo[i*10+9]['limited'][1] = 1
    parinfo[i*10+9]['limits'][1] = int(0.5*model_list[0][1]['dim'])
    parinfo[i*10+9]['fixed'] = 1
    
    parinfo[i*10+10]['value'] = model_list[i][1]['centr_offset_y']
    parinfo[i*10+10]['limited'][0] = 1
    parinfo[i*10+10]['limits'][0] = -int(0.5*model_list[0][1]['dim'])
    parinfo[i*10+10]['limited'][1] = 1
    parinfo[i*10+10]['limits'][1] = int(0.5*model_list[0][1]['dim'])
    parinfo[i*10+10]['fixed'] = 1
  
  # Input to mpfit
  functkw = {'x':model_list, 'y':boxes_avg, 'err':[boxes, box_size]}
  
  # Fit parameters
  print(model_list,boxes_avg)
  m = mpfit.mpfit(model_func, parinfo=parinfo, functkw=functkw, quiet=0)
  p = m.params
  print(p,m.status)
  
  # Create output model_list
  for i in range(len(model_list)):
    model_list[i][1]['pa'] = p[1]
    model_list[i][1]['inclination'] = p[2]
    model_list[i][1]['exp_max'] = p[3]
    model_list[i][1]['r_max'] = p[4]
    model_list[i][1]['v_max'] = p[5]
    model_list[i][1]['v_system'] = p[6]
    model_list[i][1]['v_expansion'] = p[7]
    model_list[i][1]['a_scale'] = p[8]
    model_list[i][1]['centr_offset_x'] = p[9]
    model_list[i][1]['centr_offset_y'] = p[10]

  return model_list,parinfo

  
def model_func(p, fjac=None, x=None, y=None, err=None):
  """Only used by mpfit() in fit_parameters(). Creates a model using the
      parameters to test and returns the difference between the model
      and the data in the boxes given.
      """
  # Set meaningful names
  model_list = x
  boxes_avg = N.array(y)
  boxes = err[0]
  box_size = err[1]

  print(p)
  # Create a model_list
  for i in range(len(model_list)):
    model_list[i][1]['pa'][0] = p[1]
    model_list[i][1]['inclination'][0] = p[2]
    model_list[i][1]['exp_max'][0] = p[3]
    model_list[i][1]['r_max'][0] = p[4]
    model_list[i][1]['v_max'][0] = p[5]
    model_list[i][1]['v_system'][0] = p[6]
    model_list[i][1]['v_expansion'][0] = p[7]
    model_list[i][1]['a_scale'][0] = p[8]
    model_list[i][1]['centr_offset_x'][0] = p[9]
    model_list[i][1]['centr_offset_y'][0] = p[10]

  print(model_list)
  
  # Create the model
  model = create_vf(model_list)
  # Read the average in the boxes
  model_boxes_avg = N.array(get_boxes_avg(model, boxes, box_size))
  
  status = 0
  
  # Return status and model_boxes_avg-boxes_avg
  return [status, (model_boxes_avg - boxes_avg)]

  
def get_boxes(data):
  """Used by fit_parameters(). The velocity field is shown to the user that
      chooses boxes by clicking.
      """
  # Clean the click-file
  try:
    os.remove('/tmp/MPclick.dat-nono')
  except:
    pass
  
  # Show the velocity field and connect click-events to on_click_float
  lower=data.min()-1
  upper=data.max()+1
  font = {'fontname'   : 'Courier','color'      : 'k','fontsize'   : 20}

  P.figure(num=1, figsize=(8.14, 8), dpi=80, facecolor='w', edgecolor='k')
  P.imshow(N.swapaxes(data,0,1), vmin=lower, vmax=upper, interpolation='nearest', origin='lower', aspect='preserve')
  #MP.colorbar()
  #PyCigale.setXaxis_pc(data)
  #PyCigale.setYaxis_pc(data)
  P.title(data.p['objname'] + ' - ' + 'Radial Velocity',font)
  P.axis([0,data.nx()-1,0,data.ny()-1])
  P.connect('button_press_event', PyCigale.on_click_float)
  P.show()
  
  # Read the clicks and put the coordinates in the boxes-list
  boxes = []
  for line in open('/tmp/MPclick.dat', 'r').readlines():
    line = line.split()
    boxes += [[float(line[0]), float(line[1])]]
  
  return boxes
  

def get_boxes_avg(data, boxes, box_size):
  """Used by fit_parameters(). Computes the average value in each box defined
      by boxes and box_size.
      """
  boxes_avg = N.zeros(len(boxes),dtype='Float32')
  for i in range(len(boxes)):
    p1 = N.maximum(0,int(boxes[i][0]-box_size/2))
    p2 = N.minimum((data.nx()-1),int(boxes[i][0]+box_size/2))
    p3 = N.maximum(0,int(boxes[i][1]-box_size/2))
    p4 = N.minimum((data.ny()-1),int(boxes[i][1]+box_size/2))
    boxes_avg[i]= data[p1:p2,p3:p4].mean()
  
  return boxes_avg


def create_vf(model_list):
  vf = list(range(len(model_list)))
  for i in range(len(model_list)):
    if (model_list[i][0] == 'system'):
      vf[i] = create_system_vf(model_list[i][1])
    elif (model_list[i][0] == 'linear'):
      vf[i] = create_rot_vf(model_linear, model_list[i][1])
    elif (model_list[i][0] == 'kepler'):
      vf[i] = create_rot_vf(model_kepler, model_list[i][1])
    elif (model_list[i][0] == 'pure_kepler'):
      vf[i] = create_rot_vf(model_pure_kepler, model_list[i][1])
    elif (model_list[i][0] == 'disk'):
      vf[i] = create_rot_vf(model_disk, model_list[i][1])
    elif (model_list[i][0] == 'expansion'):
      vf[i] = create_exp_vf(model_expansion, model_list[i][1])
  
  vf_total = 0
  for i in range(len(model_list)):
    vf_total = vf_total + vf[i]
  
  return PyCigale.array(vf_total)


def parameter_dict(dim=512, pa=0.0, inclination=0.0, exp_max=512, r_max=512, v_max=100, v_system=0, v_expansion=0, a_scale=1, centr_offset_x=0, centr_offset_y=0):
  """Returnes a dictionary with parameters for a model.
      
      Usage: pars = parameter_dict(dim,pa,inclination,exp_max,r_max,v_max,
        v_system,v_expansion,a_scale,centr_offset_x,centr_offset_y)
        
      pars: The dictionary with parameters
      dim..centr_offset_y: The parameters for model_list
      """
  pars = {}
  pars['dim'] = int(dim)
  pars['pa'] = pa
  pars['inclination'] = inclination
  pars['r_max'] = r_max
  pars['v_max'] = v_max
  pars['v_system'] = v_system
  pars['v_expansion'] = v_expansion
  pars['a_scale'] = a_scale
  pars['exp_max'] = exp_max
  pars['centr_offset_x'] = centr_offset_x
  pars['centr_offset_y'] = centr_offset_y

  return pars

  
def arguments_list(pars):
  """Produces two arrays, one for the radius in each point and one for the
      angle in each point.
      
      The parameters from a parameter dictionary is used to create a map 
      of the angle and radius for each point in the galaxy (in its own 
      coordinate system) viewed through the telescope.
      
      For each point in the velocity field it calculates what radius and angle
      that point corresponds to in the galaxy-coordinates (where pa
      defines the major axis).
      
      This method is from the Fortran-program.
      
      Usage r_and_phi = arguments_list(pars)
      
      pars: A parameter dictionary
      r_and_phi: [r, phi], Two 2D-arrays with the transformed coordinates.
      """
  # Coordinates for the galaxy centrum
  centr_x = pars['dim']/2 +pars['centr_offset_x']
  centr_y = pars['dim']/2 +pars['centr_offset_y']
  
  # The empty maps
  r = N.zeros((pars['dim'],pars['dim']), dtype='Float64')
  phi = N.zeros((pars['dim'],pars['dim']), dtype='Float64')
   
  # Transform each point
  for x in range(pars['dim']):
    for y in range(pars['dim']):
      temp_x = x-centr_x
      temp_y = y-centr_y
      
      pa = (-pars['pa']-90)*M.pi/180
      inclination = pars['inclination']*M.pi/180
      
      x_corr = temp_x*N.cos(pa) - temp_y*N.sin(pa)
      y_corr = temp_x*N.sin(pa) + temp_y*N.cos(pa)
      
      y_corr = y_corr*N.cos(inclination)
      
      phi[x,y] = N.arctan2(y_corr, x_corr)
      r[x,y] = N.sqrt(x_corr**2+y_corr**2)

  r_and_phi = [r, phi]
  
  return r_and_phi

  
def create_exp_vf(model, pars):
  """Creates a VF one of the model-functions (model_expansion is the only one 
      available atm.) with the parameter dictionary 'pars'.
      
      Usage: vf_exp = create_expansion_vf(model, pars)
      
      pars: The parameter dictionary to use
      vf_exp: A velocity field (note: this field is just a pure numarray, so no 
        p-list). If you want to keep the p-list from 'vf' you can do this:
        vf_exp = vf.copy()
        vf_exp[:,:] = create_expansion_vf(pars)
  """
  r_and_phi = arguments_list(pars)

  arguments = [r_and_phi[0], pars]
  v_model_exp = model_expansion(*arguments)

  vf = v_model_exp*N.sin(r_and_phi[1])*N.sin(pars['inclination'][0]*M.pi/180)
   
  return vf

  
def create_rot_vf(model, pars):
  """Creates a VF using one of the model-functions (model_disk, model_kepler, 
      model_linear, model_pure_kepler) with the parameter dictionary 'pars'.
      
      Usage: vf_rot = create_rot_vf(model, pars)
      
      model: One of the functions listed above
      pars: The parameter dictionary to use
      vf_rot: A velocity field (note: this field is just a pure numarray, so no 
        p-list). If you want to keep the p-list from 'vf' you can do this:
        vf_rot = vf.copy()
        vf_rot[:,:] = create_rot_vf(model, pars)
  """
  r_and_phi = arguments_list(pars)

  arguments = [r_and_phi[0], pars]
  # Call model-function
  v_model_rot = model(*arguments)

  # Projection to the rotation plane
  vf = v_model_rot*N.cos(r_and_phi[1])*N.sin(pars['inclination']*M.pi/180)
  
  return vf

  
def create_system_vf(pars):
  """Creates a VF with a system-velocity using the parameter dictionary 'pars'.
      
      Usage: vf_sys = create_system_vf(pars)
      
      pars: The parameter dictionary to use
      vf_sys: A velocity field (note: this field is just a pure numarray, so no 
        p-list). If you want to keep the p-list from 'vf' you can do this:
        vf_sys = vf.copy()
        vf_sys[:,:] = create_system_vf(pars)
  """
  #print pars['v_system']
  vf = N.zeros((pars['dim'],pars['dim']), dtype='Float64') + pars['v_system']
  return vf
  
  
def model_expansion(r, pars):
  """Creates an expansion velocity field. The velocity is v_max in the centrum
      and decreases as 1/r**2, stretched such that it is 0.1*v_max at the radius 
      exp_max.
      
      Usage: arr = model_expansion(r, pars)
      
      r: An array with radii in every point, preferably from argument_list()
      pars: The parameter dictionary to use
      arr: An array with velocities in every point
      """
  # Bugfix, we don't want to change r, just a copy of it
  arr = r.copy()
  
  # Set variables
  len_x = arr.shape[0]
  len_y = arr.shape[1]
  
  v_max = pars['v_expansion'][0]
  exp_max = pars['exp_max'][0]

  # Special case if exp_max=0
  if (exp_max == 0):
    arr[:,:] = 0
    return arr
  
  #Calculate the velocity in ervery point
  arr.setshape((len_x*len_y))

  arr = v_max/(( ((N.sqrt(10)-1)/exp_max)*arr+1)**2)
  
  arr.setshape((len_x, len_y)) 
  return arr

  
def model_disk(r, pars):
  """Creates an disk rotation velocity field. 
      
      Usage: arr = model_disk(r, pars)
      
      r: An array with radii in every point, preferably from argument_list()
      pars: The parameter dictionary to use
      arr: An array with velocities in every point
      """
  # Bugfix, we don't want to change r, just a copy of it
  arr = r.copy()
  
  # Set variables
  len_x = arr.shape[0]
  len_y = arr.shape[1]
  
  a = pars['a_scale']
  vm = pars['v_max']

  # Calculate the velocity in ervery point
  arr.setshape((len_x*len_y))
  
  t1 = N.where(arr==0)
  t2 = N.where(arr>0)
  
  arr[t1] = 0.0
  arr[t2] = N.sqrt(N.fabs(vm**2 * ( 1 - ( N.arctan2(arr[t2],a) * (a/arr[t2]) ) ) ) )
  
  arr.setshape((len_x, len_y)) 
  
  return arr
  
  
def model_kepler(r, pars):
  """Creates an kepler rotation velocity field. Linear up to v_max*r/r_max and 
      then decreasing as v_max*sqrt(r_max/r).
      
      Usage: arr = model_disk(r, pars)
      
      r: An array with radii in every point, preferably from argument_list()
      pars: The parameter dictionary to use
      arr: An array with velocities in every point
      """
  # Bugfix, we don't want to change r, just a copy of it
  arr = r.copy()
  
  # Set variables
  len_x = arr.shape[0]
  len_y = arr.shape[1]
  
  rm = pars['r_max'][0]
  vm = pars['v_max'][0]
  
  # Special case when r_max = 0
  if (rm == 0):
    arr[:,:] = 0
    return arr
  
  # Calculate the velocity in ervery point
  arr.setshape((len_x*len_y))
  
  t1 = N.where(arr<rm)
  t2 = N.where(arr>=rm)
  
  arr[t1] = vm*arr[t1]/rm
  arr[t2] = vm*N.sqrt(rm/arr[t2])
  
  arr.setshape((len_x, len_y))
  return arr
  

def model_pure_kepler(r, pars):
  """Creates an pure kepler rotation velocity field. Decreasing as 
      v_max*sqrt(r_max/r). The velocity is set to 0 in the centrum.
      
      Usage: arr = model_disk(r, pars)
      
      r: An array with radii in every point, preferably from argument_list()
      pars: The parameter dictionary to use
      arr: An array with velocities in every point
      """
  # Bugfix, we don't want to change r, just a copy of it
  arr = r.copy()
  
  # Set variables
  len_x = arr.shape[0]
  len_y = arr.shape[1]
  
  rm = pars['r_max'][0]
  vm = pars['v_max'][0]
  
  # Special case when r_max = 0
  if (rm == 0):
    arr[:,:] = 0
    return arr
  
  # Calculate the velocity in ervery point
  arr.setshape((len_x*len_y))
  
  t1 = N.where(arr<0.5)
  t2 = N.where(arr>=0.5)
  
  arr[t1] = 0.0
  arr[t2] = vm*N.sqrt(rm/arr[t2])
  
  arr.setshape((len_x, len_y))
  return arr
  
  
def model_linear(r, pars):
  """Creates an linear rotation velocity field. Increasing as v_max*r/r_max,
      up to r_max, then it is held fix (at v_max).
      
      Usage: arr = model_disk(r, pars)
      
      r: An array with radii in every point, preferably from argument_list()
      pars: The parameter dictionary to use
      arr: An array with velocities in every point
      """
  # Bugfix, we don't want to change r, just a copy of it
  arr = r.copy()
  
  # Set variables
  len_x = arr.shape[0]
  len_y = arr.shape[1]
  
  rm = pars['r_max']
  vm = pars['v_max']

  # Calculate the velocity in ervery point
  arr.setshape((len_x*len_y))
  
  t1 = N.where(arr<rm)
  t2 = N.where(arr>=rm)
  
  arr[t1] = vm*arr[t1]/rm
  arr[t2] = vm
  
  arr.setshape((len_x, len_y))
  return arr
 


 
###############################################
# The old stuff
###############################################

def model(model,name,*args):
  
  if model == 'linear':
    print("creating the linear field...")
    result,err=create_linear (args)
    if err != 0:
      print("error")
    else:  
      print("done.")
    
    
  elif model == 'disk':
    print("creating the disk...")
    result,err=create_disk (args)
    if err != 0:
      print("error")
    else:  
      print("done.")
    
  elif model == 'shell':
    print("creating the shell...")
    result,err=create_shell (args)
    if err != 0:
      print("error")
    else:  
      print("done.")

  elif model == 'sphere':
    print("creating the sphere...")
    result,err=create_sphere (args)
    if err != 0:
      print("error")
    else:  
      print("done.")
    
  else:
    print("unknown model")
    result,err=N.zeros((1,1)),1

  # write the file if all went fine
  if err == 0:
    os.popen('rm ' + name)
    IO.write_fits(result, name)

  else:
    print("no file written")



def create_linear (arg):
  """
  makes a velocity field with a simple gradient.
  rotation possible.
  """
  if len(arg) != 3:
    print("Wrong number of arguments")
    return N.zeros((1,1)),1
  else:
    dim=arg[0]             # dimension of the file.
    vel=float(arg[1])      # peak velocity at r=dim (it becomes bigger than
                           # that in the corners if you rotate of course).
    rot=M.radians(arg[2])  # rotation from top to left in degrees.
    
  vf = N.zeros((dim,dim), dtype='float32')
  cent = dim / 2.
  diff = vel / cent
  normvec=[M.sin(rot),M.cos(rot)]
  
  for i in range(dim):
    icent=i-cent
    for j in range(dim):

      vf[i,j]= diff * N.dot(normvec,[icent,j-cent])

      
  return vf,0



def create_disk (arg):
  """
  creates a velocity field of a thin disk galaxy
  according to Sparke, Gallagher's
  "Galaxies in the Universe", equations 2.20 and 5.4 .
  Parameters taken from there as well.
  """
  if len(arg) != 5:
    print("Wrong number of arguments")
    return N.zeros((1,1)),1
  else:
    dim=arg[0]                # dimension of the file.
    vel=float(arg[1])         # max velocity
    a=float(arg[2])           # scale parameter
    incl=M.radians(arg[3])    # inclination
    rot=M.radians(arg[4])     # rotation from top to left in degrees.
    
  vf = N.zeros((dim,dim), dtype='float32')
  cent = dim / 2.
  minaxis=[M.sin(rot),M.cos(rot)]
  majaxis=[M.cos(rot),-M.sin(rot)]


  for i in range(dim):
    icent=i-cent
    for j in range(dim):
      
      d_min=N.dot(majaxis,[icent,j-cent])
      d_maj=N.dot(minaxis,[icent,j-cent]) / N.cos(incl)
      
      r = N.hypot(d_min, d_maj)
      if r == 0.0:
        vf[i,j]=0

      else:  
        
        v_r = v_r_disk(vel,a,r)

        vf[i,j]= v_r * N.sin(incl) * (d_min / r) # note that (d_min / r) = cos Phi
        

  return vf,0    


def create_shell(arg):
  """
  Velocity field of a rotating (opt thick) ball, looking at the surface
  Parameter vel is the velocity IN the image plane at the center.
  R is the radius of the sphere which need not be the same as the
  image size.
  rotation means the rot-axis from top to left in degrees.
  """
  if len(arg) != 4:
    print("Wrong number of arguments")
    return N.zeros((1,1)),1
  else:
    dim=arg[0]                # dimension of the file.
    R=float(arg[1])           # Radius (see above)
    vel=float(arg[2])         # vel. (see above)
    rot=M.radians(arg[3])     # rotation from top to left in degrees.
    
  vf = N.zeros((dim,dim), dtype='float32')
  cent = dim / 2.
  perpaxis=[M.sin(rot),M.cos(rot)]
  rotaxis=[M.cos(rot),-M.sin(rot)] 
  omega=vel / R


  for i in range(dim):
    icent=i-cent
    for j in range(dim):
      
      d_rot=N.dot(perpaxis,[icent,j-cent]) # distance from rotation axis
      d_cent=N.dot(rotaxis,[icent,j-cent]) # distance from center along rot-axis
      d_centS = d_cent / R * M.pi / 2.
      r = N.hypot(d_rot, d_cent)
      if r == 0:
        vf[i,j]= 0
      elif r >= R:
        vf[i,j]= 0
      else:
        vf[i,j]= d_rot * omega
        #r_ring= N.sqrt(R**2 - d_cent**2) # radius of the circle a height d_cent
        #alpha = N.arcsin(d_rot / r_ring)
        #vf[i,j]= N.sin(alpha) * (omega * r_ring)
        #N.sqrt((omega*d_rot)**2 * ( 1 - (d_rot / r_ring)**2 ) )
        #v_r_ring(r_ring * omega,d_rot,r)
        
  return vf,0    
  


def create_sphere(arg):
  """
  Velocity field of a rotating opt. thin ball, looking through it.
  This will be more tricky and probably needs more assumtions.
  """
  return N.zeros((1,1)),1


  

def v_r_disk(vel,a,r):
  """
  V(r) according to eq. 2.20 from Sparke, Gallagher
  """
  return  N.sqrt(N.fabs(vel**2 * ( 1 - ( N.arctan2(r,a) * (a/r) ) ) ) ) 



def get_box(vf,x):
  """
  gives the average in predefined boxes
  """

  if x == 1:
    aver=mean(mean(vf[20:30,100:110]))
  elif x==2:
    aver=mean(mean(vf[70:80,110:120]))
  elif x==3:
    aver=mean(mean(vf[95:105,95:105]))
  elif x==4:
    aver=mean(mean(vf[105:115,65:75]))
  elif x==5:
    aver=mean(mean(vf[95:105,25:35]))
  elif x==6:
    aver=mean(mean(vf[50:60,15:25]))
  elif x==7:
    aver=mean(mean(vf[20:30,40:50]))
  elif x==8:
    aver=mean(mean(vf[20:30,70:80]))
  elif x==9:
    aver=mean(mean(vf[20:30,100:110]))
  elif x==10:
    aver=mean(mean(vf[20:30,100:110]))
  else:
    aver=0

  return aver


def modelfu_linear(param,x):
  """
  gives back the model velocity at the desired points
  """

  dim=64.

  if x == 1:
    X=25
    Y=105
  elif x==2:
    X=75
    Y=115
  elif x==3:
    X=100
    Y=100
  elif x==4:
    X=110
    Y=70
  elif x==5:
    X=100
    Y=30
  elif x==6:
    X=55
    Y=20
  elif x==7:
    X=25
    Y=45
  elif x==8:
    X=25
    Y=75
  else:
    return 1

  rot=param[1] / 360 * 2 * M.pi
  #print param[0], rot, rot / 2 / M.pi * 360, X, Y
  return (param[0]/dim) * Nu.dot([Nu.sin(rot),Nu.cos(rot)],[X-dim,Y-dim])



def fit_linear(name):
  """
  testing a least square fit
  """



  rawdata=IO.read_fits(name)

  data= [ (1, get_box(rawdata,1)),
          (2, get_box(rawdata,2)),
          (3, get_box(rawdata,3)),
          (4, get_box(rawdata,4)),
          (5, get_box(rawdata,5)),
          (6, get_box(rawdata,6)),
          (7, get_box(rawdata,7)),
          (8, get_box(rawdata,8)) ]
  
  result,chi = LS.leastSquaresFit(modelfu_linear, (100,-30,), data)
  print(result[0], M.fmod(result[1],360), chi)



def fit_test():
  """
  testing a least square fit
  """

  def modelfu_test(param,t):
    return param[0]*t

  data= [ (1, 1),
          (2, 2),
          (3, 3),
          (4, 4),
          (5, 5),
          (6, 6),
          (7, 7),
          (8, 8) ]
  
  print(LS.leastSquaresFit(modelfu_test, (2,), data))