Scripts/ExploreCTF: apply_ctf.py

#!/usr/bin/env giplt
# Tab-Size: 4 (soft)

import getopt, sys
import ex.alg

def setup_ctf(**kwargs):
    """
    Returns a CTF object with your parameters of choice. The convention for
    defocus is such that positive values depict underfocus and negative
    values are overfocus.
    """
    options = {
        'defocus'   : 858,
        'cs'        : 2e7,
        'cc'        : 2e7
    }
    options.update(**kwargs)
    microscope_data=ex.MicroscopeData(options['cs'],options['cc'],0)
    microscope_data.SetAccelerationVoltage(100*1.0e23)
    ctf=ex.alg.CTF(ex.TCIFData(microscope_data,
        ex.Defocus(options['defocus'])))
    return ctf
    
def ctf_values(**kwargs):
    """
    Returns a one-dimensional array containing the CTF values.
    """
    options = {
        'sampling'  : 1,
    }
    options.update(kwargs)
    ctf=setup_ctf(**kwargs)
    im=CreateImage(Size(2048,1))
    im.SetPixelSampling(Vec3(options['sampling'],
        options['sampling'],options['sampling']))
    im.ApplyIP(alg.FFT())   
    im.ApplyIP(alg.Fill(1))        
    ctf_values=im.Apply(ctf)
    return [ctf_values.GetComplex(Point(x,0)).real for x in range(0,1024)]

def plot_ctf(**kwargs):
    """
    Returns plot data for CTF with the given parameters
    """
    ctf_val=ctf_values(**kwargs)
    pd=gui.PlotData()
    for x in range(0,1024):
        pd.AddXY(x,ctf_val[x])
    return pd

    
def print_usage():
    print 'Usage: apply_ctf.py <options> image'
    print 'Options:'
    print '   --defocus=<value>    Defocus in A. Positive values for underfocus'
    print '         -cc=<value>    Chromatic abberation in mm'
    print '         -cs=<value>    Spherical abberation in mm'
    print '  --sampling=<value>    Sampling for image in A/px'
    print ''
    print 'values default to defocus=700A, cc=2mm, cs=2mm, sampling=1A/px'

def image_ctf(image,**kwargs):
    ctf=setup_ctf(**kwargs)
    image=LoadImage(image)
    image.SetPixelSampling(Vec3(1,1,1))
    image.ApplyIP(alg.FFT())
    ctf_mod=image.Apply(ctf)
    return ctf_mod.Apply(alg.FFT())
    
    
def tuple_to_dict(t):
    d = {}
    for k,v in t:
        d[k[2:]] = int(v)
    return d

def get_options():
    options=['cs=','cc=','defocus=','sampling=']
    try:
        opts,args=getopt.getopt(sys.argv[1:], '', options)
    except getopt.GetOptError, err:
        print str(err)
        print_usage()
        sys.exit(0)
    if len(args) != 1:
        print_usage()
        sys.exit(0)
    default= {
        'defocus'   : 700,
        'cs'        : 2,
        'cc'        : 2,
        'sampling'  : 1
    }
    default.update(tuple_to_dict(opts))
    default['cs']*=1e7
    default['cc']*=1e7
    return default,args[0]
        
opts,image_file=get_options()

print opts

pv=gui.CreatePlotViewer()

pv.AddData(plot_ctf(**opts)).SetMode(pv.LINES)

im=image_ctf(image_file,**opts)
v=Viewer(im)
fft=im.Apply(alg.FFT())
v_fft=Viewer(fft)