source: sasview/sanscalculator/src/sans/calculator/sans_gen.py @ 3b148c3

ESS_GUIESS_GUI_DocsESS_GUI_batch_fittingESS_GUI_bumps_abstractionESS_GUI_iss1116ESS_GUI_iss879ESS_GUI_iss959ESS_GUI_openclESS_GUI_orderingESS_GUI_sync_sascalccostrafo411magnetic_scattrelease-4.1.1release-4.1.2release-4.2.2release_4.0.1ticket-1009ticket-1094-headlessticket-1242-2d-resolutionticket-1243ticket-1249ticket885unittest-saveload
Last change on this file since 3b148c3 was 968f72d, checked in by Jae Cho <jhjcho@…>, 12 years ago

GENSANS: set all Gen files as default in fileDialog.

  • Property mode set to 100644
File size: 43.0 KB
Line 
1"""
2SANS generic computation and sld file readers
3"""
4from sans.models.BaseComponent import BaseComponent
5import sans.models.sans_extension.sld2i as mod
6from periodictable import formula
7from periodictable import nsf
8import numpy
9import os
10import copy
11
12MFactor_AM = 2.853E-12
13MFactor_mT = 2.3164E-9
14METER2ANG = 1.0E+10
15#Avogadro constant [1/mol]
16NA = 6.02214129e+23
17
18def mag2sld(mag, v_unit=None):
19    """
20    Convert magnetization to magnatic SLD
21    sldm = Dm * mag where Dm = gamma * classical elec. radius/(2*Bohr magneton)
22    Dm ~ 2.853E-12 [A^(-2)] ==> Shouldn't be 2.90636E-12 [A^(-2)]???
23    """ 
24    if v_unit == "A/m":
25        factor = MFactor_AM
26    elif v_unit == "mT":
27        factor = MFactor_mT
28    else:
29        raise ValueError, "Invalid valueunit"
30    sld_m = factor * mag
31    return sld_m
32
33def transform_center(pos_x, pos_y, pos_z):
34    """
35    re-center
36   
37    :return: posx, posy, posz   [arrays]
38    """ 
39    posx = pos_x - (min(pos_x) + max(pos_x)) / 2.0
40    posy = pos_y - (min(pos_y) + max(pos_y)) / 2.0
41    posz = pos_z - (min(pos_z) + max(pos_z)) / 2.0 
42    return posx, posy, posz   
43
44class GenSAS(BaseComponent):
45    """
46    Generic SAS computation Model based on sld (n & m) arrays
47    """
48    def __init__(self):
49        """
50        Init
51       
52        :Params sld_data: MagSLD object
53        """
54        # Initialize BaseComponent
55        BaseComponent.__init__(self)
56        self.sld_data = None
57        self.data_pos_unit = None
58        self.data_x = None
59        self.data_y = None
60        self.data_z = None
61        self.data_sldn = None
62        self.data_mx = None
63        self.data_my = None
64        self.data_mz = None
65        self.data_vol = None #[A^3]
66        self.is_avg = False
67        ## Name of the model
68        self.name = "GenSAS"
69        ## Define parameters
70        self.params = {}
71        self.params['scale']       = 1.0
72        self.params['background']  = 0.0
73        self.params['solvent_SLD']     = 0.0
74        self.params['total_volume'] = 1.0
75        self.params['Up_frac_in']     = 1.0
76        self.params['Up_frac_out']    = 1.0
77        self.params['Up_theta']  = 0.0
78        self.description = 'GenSAS'
79        ## Parameter details [units, min, max]
80        self.details = {}
81        self.details['scale']      = ['', None, None]
82        self.details['background'] = ['[1/cm]', None, None]
83        self.details['solvent_SLD']    = ['1/A^(2)', None, None]
84        self.details['total_volume']    = ['A^(3)', None, None]
85        self.details['Up_frac_in']    = ['[u/(u+d)]', None, None]
86        self.details['Up_frac_out']   = ['[u/(u+d)]', None, None]
87        self.details['Up_theta'] = ['[deg]', None, None]
88        # fixed parameters
89        self.fixed = []
90       
91    def set_pixel_volumes(self, volume):     
92        """
93        Set the volume of a pixel in (A^3) unit
94        :Param volume: pixel volume [float]
95        """
96        if self.data_vol == None:
97            raise
98        self.data_vol = volume
99   
100    def set_is_avg(self, is_avg=False): 
101        """
102        Sets is_avg: [bool]
103        """
104        self.is_avg = is_avg
105         
106    def _gen(self, x, y, i):
107        """
108        Evaluate the function
109        :Param x: array of x-values
110        :Param y: array of y-values
111        :Param i: array of initial i-value
112        :return: function value
113        """
114        pos_x = self.data_x
115        pos_y = self.data_y
116        pos_z = self.data_z
117        len_x = len(pos_x)
118        if self.is_avg == None:
119            len_x *= -1
120            pos_x, pos_y, pos_z = transform_center(pos_x, pos_y, pos_z)
121        len_q = len(x)
122        sldn = copy.deepcopy(self.data_sldn)
123        sldn -= self.params['solvent_SLD']
124        model = mod.new_GenI(len_x, pos_x, pos_y, pos_z, 
125                             sldn, self.data_mx, self.data_my, 
126                             self.data_mz, self.data_vol,
127                             self.params['Up_frac_in'], 
128                             self.params['Up_frac_out'], 
129                             self.params['Up_theta'])
130        if y == []:
131            mod.genicom(model, len_q, x, i)
132        else:
133            mod.genicomXY(model, len_q, x, y, i)
134        vol_correction = self.data_total_volume / self.params['total_volume']
135        return  self.params['scale'] * vol_correction * i + \
136                        self.params['background']
137       
138    def set_sld_data(self, sld_data=None):   
139        """
140        Sets sld_data
141        """
142        self.sld_data = sld_data
143        self.data_pos_unit = sld_data.pos_unit
144        self.data_x = sld_data.pos_x
145        self.data_y = sld_data.pos_y
146        self.data_z = sld_data.pos_z
147        self.data_sldn = sld_data.sld_n
148        self.data_mx = sld_data.sld_mx
149        self.data_my = sld_data.sld_my
150        self.data_mz = sld_data.sld_mz
151        self.data_vol = sld_data.vol_pix
152        self.data_total_volume = sum(sld_data.vol_pix)
153        self.params['total_volume'] = sum(sld_data.vol_pix)
154       
155    def getProfile(self):
156        """
157        Get SLD profile
158       
159        : return: sld_data
160        """
161        return self.sld_data
162     
163    def run(self, x = 0.0):
164        """
165        Evaluate the model
166        :param x: simple value
167        :return: (I value)
168        """
169        if x.__class__.__name__ == 'list':
170            if len(x[1]) > 0:
171                msg = "Not a 1D."
172                raise ValueError, msg
173            i_out = numpy.zeros_like(x[0])
174            # 1D I is found at y =0 in the 2D pattern
175            out = self._gen(x[0], [], i_out )
176            return out
177        else:
178            msg = "Q must be given as list of qx's and qy's"
179            raise ValueError, msg
180   
181    def runXY(self, x = 0.0):
182        """
183        Evaluate the model
184        :param x: simple value
185        :return: I value
186        :Use this runXY() for the computation
187        """
188        if x.__class__.__name__ == 'list':
189            i_out = numpy.zeros_like(x[0])
190            out = self._gen(x[0], x[1], i_out)
191            return out
192        else:
193            msg = "Q must be given as list of qx's and qy's"
194            raise ValueError, msg
195       
196    def evalDistribution(self, qdist):
197        """
198        Evaluate a distribution of q-values.
199        * For 1D, a numpy array is expected as input:
200            evalDistribution(q)   
201        where q is a numpy array.
202        * For 2D, a list of numpy arrays are expected: [qx_prime,qy_prime],
203          where 1D arrays, 
204        :param qdist: ndarray of scalar q-values or list [qx,qy]
205                    where qx,qy are 1D ndarrays
206        """
207        if qdist.__class__.__name__ == 'list':
208            if len(qdist[1]) < 1:
209                out = self.run(qdist)
210            else:
211                out = self.runXY(qdist)
212            return out
213        else:
214            mesg = "evalDistribution is expecting an ndarray of "
215            mesg += "a list [qx,qy] where qx,qy are arrays."
216            raise RuntimeError, mesg
217     
218class OMF2SLD:
219    """
220    Convert OMFData to MAgData
221    """
222    def __init__(self): 
223        """
224        Init
225        """
226        self.pos_x = None
227        self.pos_y = None
228        self.pos_z = None
229        self.mx = None
230        self.my = None
231        self.mz = None
232        self.sld_n = None
233        self.vol_pix = None
234        self.output = None
235        self.omfdata = None
236   
237    def set_data(self, omfdata, shape='rectangular'):
238        """
239        Set all data
240        """
241        self.omfdata = omfdata
242        length = int(omfdata.xnodes * omfdata.ynodes * omfdata.znodes)
243        pos_x = numpy.arange(omfdata.xmin, 
244                             omfdata.xnodes*omfdata.xstepsize + omfdata.xmin, 
245                             omfdata.xstepsize)
246        pos_y = numpy.arange(omfdata.ymin, 
247                             omfdata.ynodes*omfdata.ystepsize + omfdata.ymin, 
248                             omfdata.ystepsize)
249        pos_z = numpy.arange(omfdata.zmin, 
250                             omfdata.znodes*omfdata.zstepsize + omfdata.zmin, 
251                             omfdata.zstepsize)
252        self.pos_x = numpy.tile(pos_x, int(omfdata.ynodes * omfdata.znodes))
253        self.pos_y = pos_y.repeat(int(omfdata.xnodes))
254        self.pos_y = numpy.tile(self.pos_y, int(omfdata.znodes))
255        self.pos_z = pos_z.repeat(int(omfdata.xnodes * omfdata.ynodes))
256        self.mx = omfdata.mx
257        self.my = omfdata.my
258        self.mz = omfdata.mz
259        self.sld_n = numpy.zeros(length)
260           
261        if omfdata.mx == None:
262            self.mx = numpy.zeros(length)
263        if omfdata.my == None:
264            self.my = numpy.zeros(length)
265        if omfdata.mz == None:
266            self.mz = numpy.zeros(length)
267       
268        self._check_data_length(length)
269        self.remove_null_points(False, False)
270        mask = numpy.ones(len(self.sld_n), dtype=bool)
271        if shape.lower() == 'ellipsoid':
272            try:
273                # Pixel (step) size included
274                x_c = max(self.pos_x) + min(self.pos_x)
275                y_c = max(self.pos_y) + min(self.pos_y)
276                z_c = max(self.pos_z) + min(self.pos_z)
277                x_d = max(self.pos_x) - min(self.pos_x)
278                y_d = max(self.pos_y) - min(self.pos_y)
279                z_d = max(self.pos_z) - min(self.pos_z)
280                x_r = (x_d + omfdata.xstepsize) / 2.0
281                y_r = (y_d + omfdata.ystepsize) / 2.0
282                z_r = (z_d + omfdata.zstepsize) / 2.0
283                x_dir2 = ((self.pos_x - x_c / 2.0) / x_r)
284                x_dir2 *= x_dir2
285                y_dir2 = ((self.pos_y - y_c / 2.0) / y_r)
286                y_dir2 *= y_dir2
287                z_dir2 = ((self.pos_z - z_c / 2.0) / z_r)
288                z_dir2 *= z_dir2
289                mask = (x_dir2 + y_dir2 + z_dir2) <= 1.0
290            except:
291                pass
292        self.output = MagSLD(self.pos_x[mask], self.pos_y[mask], 
293                             self.pos_z[mask], self.sld_n[mask], 
294                             self.mx[mask], self.my[mask], self.mz[mask])
295        self.output.set_pix_type('pixel')
296        self.output.set_pixel_symbols('pixel')
297       
298    def get_omfdata(self):
299        """
300        Return all data
301        """
302        return self.omfdata
303   
304    def get_output(self):
305        """
306        Return output
307        """
308        return self.output
309   
310    def _check_data_length(self, length):
311        """
312        Check if the data lengths are consistent
313        :Params length: data length
314        """
315        msg = "Error: Inconsistent data length."
316        if len(self.pos_x) != length:
317            raise ValueError, msg
318        if len(self.pos_y) != length:
319            raise ValueError, msg
320        if len(self.pos_z) != length:
321            raise ValueError, msg
322        if len(self.mx) != length:
323            raise ValueError, msg
324        if len(self.my) != length:
325            raise ValueError, msg
326        if len(self.mz) != length:
327            raise ValueError, msg
328       
329    def remove_null_points(self, remove=False, recenter=False):
330        """
331        Removes any mx, my, and mz = 0 points
332        """
333        if remove:
334            is_nonzero = (numpy.fabs(self.mx) + numpy.fabs(self.my) + 
335                          numpy.fabs(self.mz)).nonzero() 
336            if len(is_nonzero[0]) > 0: 
337                self.pos_x = self.pos_x[is_nonzero]
338                self.pos_y = self.pos_y[is_nonzero]
339                self.pos_z = self.pos_z[is_nonzero]
340                self.sld_n = self.sld_n[is_nonzero]
341                self.mx = self.mx[is_nonzero]
342                self.my = self.my[is_nonzero]
343                self.mz = self.mz[is_nonzero]
344        if recenter:
345            self.pos_x -= (min(self.pos_x) + max(self.pos_x)) / 2.0
346            self.pos_y -= (min(self.pos_y) + max(self.pos_y)) / 2.0
347            self.pos_z -= (min(self.pos_z) + max(self.pos_z)) / 2.0
348 
349    def get_magsld(self):
350        """
351        return MagSLD
352        """
353        return self.output
354   
355       
356class OMFReader:
357    """
358    Class to load omf/ascii files (3 columns w/header).
359    """
360    ## File type
361    type_name = "OMF ASCII"
362   
363    ## Wildcards
364    type = ["OMF files (*.OMF, *.omf)|*.omf"]
365    ## List of allowed extensions
366    ext = ['.omf', '.OMF']
367       
368    def read(self, path):
369        """
370        Load data file
371        :param path: file path
372        :return: x, y, z, sld_n, sld_mx, sld_my, sld_mz
373        """
374        #data_conv_m = None
375        desc = ""
376        mx = numpy.zeros(0)
377        my = numpy.zeros(0)
378        mz = numpy.zeros(0)
379        try:
380            data_conv_val = None
381            data_conv_mesh = None
382            input_f = open(path, 'rb')
383            buff = input_f.read()
384            lines = buff.split('\n')
385            input_f.close()
386            output = OMFData()
387            valueunit = None
388            for line in lines:
389                toks = line.split()
390                # Read data
391                try:
392                    _mx = float(toks[0])
393                    _my = float(toks[1])
394                    _mz = float(toks[2])
395                    _mx = mag2sld(_mx, valueunit)
396                    _my = mag2sld(_my, valueunit)
397                    _mz = mag2sld(_mz, valueunit)
398                    mx = numpy.append(mx, _mx)
399                    my = numpy.append(my, _my)
400                    mz = numpy.append(mz, _mz)
401                except:
402                    # Skip non-data lines
403                    pass
404                #Reading Header; Segment count ignored
405                s_line = line.split(":", 1)
406                if s_line[0].lower().count("oommf") > 0:
407                    oommf = s_line[1].lstrip()
408                if s_line[0].lower().count("title") > 0:
409                    title = s_line[1].lstrip()
410                if s_line[0].lower().count("desc") > 0:
411                    desc += s_line[1].lstrip()
412                    desc += '\n'
413                if s_line[0].lower().count("meshtype") > 0:
414                    meshtype = s_line[1].lstrip()
415                if s_line[0].lower().count("meshunit") > 0:
416                    meshunit = s_line[1].lstrip()
417                    if meshunit.count("m") < 1:
418                        msg = "Error: \n"
419                        msg += "We accept only m as meshunit"
420                        raise ValueError, msg
421                if s_line[0].lower().count("xbase") > 0:
422                    xbase = s_line[1].lstrip()
423                if s_line[0].lower().count("ybase") > 0:
424                    ybase = s_line[1].lstrip()
425                if s_line[0].lower().count("zbase") > 0:
426                    zbase = s_line[1].lstrip()
427                if s_line[0].lower().count("xstepsize") > 0:
428                    xstepsize = s_line[1].lstrip()
429                if s_line[0].lower().count("ystepsize") > 0:
430                    ystepsize = s_line[1].lstrip()
431                if s_line[0].lower().count("zstepsize") > 0:
432                    zstepsize = s_line[1].lstrip()
433                if s_line[0].lower().count("xnodes") > 0:
434                    xnodes = s_line[1].lstrip()
435                if s_line[0].lower().count("ynodes") > 0:
436                    ynodes = s_line[1].lstrip()
437                if s_line[0].lower().count("znodes") > 0:
438                    znodes = s_line[1].lstrip()
439                if s_line[0].lower().count("xmin") > 0:
440                    xmin = s_line[1].lstrip()
441                if s_line[0].lower().count("ymin") > 0:
442                    ymin = s_line[1].lstrip()
443                if s_line[0].lower().count("zmin") > 0:
444                    zmin = s_line[1].lstrip()
445                if s_line[0].lower().count("xmax") > 0:
446                    xmax = s_line[1].lstrip()
447                if s_line[0].lower().count("ymax") > 0:
448                    ymax = s_line[1].lstrip()
449                if s_line[0].lower().count("zmax") > 0:
450                    zmax = s_line[1].lstrip()
451                if s_line[0].lower().count("valueunit") > 0:
452                    valueunit = s_line[1].lstrip().rstrip()
453                if s_line[0].lower().count("valuemultiplier") > 0:
454                    valuemultiplier = s_line[1].lstrip()
455                if s_line[0].lower().count("valuerangeminmag") > 0:
456                    valuerangeminmag = s_line[1].lstrip()
457                if s_line[0].lower().count("valuerangemaxmag") > 0:
458                    valuerangemaxmag = s_line[1].lstrip()
459                if s_line[0].lower().count("end") > 0:
460                    #output.set_sldms(mx, my, mz)   
461                    output.filename = os.path.basename(path)
462                    output.oommf = oommf
463                    output.title = title
464                    output.desc = desc
465                    output.meshtype = meshtype
466                    output.xbase = float(xbase) * METER2ANG
467                    output.ybase = float(ybase) * METER2ANG
468                    output.zbase = float(zbase) * METER2ANG
469                    output.xstepsize = float(xstepsize) * METER2ANG
470                    output.ystepsize = float(ystepsize) * METER2ANG
471                    output.zstepsize = float(zstepsize) * METER2ANG
472                    output.xnodes = float(xnodes)
473                    output.ynodes = float(ynodes)
474                    output.znodes = float(znodes)
475                    output.xmin = float(xmin) * METER2ANG
476                    output.ymin = float(ymin) * METER2ANG
477                    output.zmin = float(zmin) * METER2ANG
478                    output.xmax = float(xmax) * METER2ANG
479                    output.ymax = float(ymax) * METER2ANG
480                    output.zmax = float(zmax) * METER2ANG
481                    output.valuemultiplier = valuemultiplier
482                    output.valuerangeminmag = mag2sld(float(valuerangeminmag), \
483                                                      valueunit)
484                    output.valuerangemaxmag = mag2sld(float(valuerangemaxmag), \
485                                                      valueunit)
486            output.set_m(mx, my, mz)
487            return output
488        except:
489            msg = "%s is not supported: \n" % path
490            msg += "We accept only Text format OMF file."
491            raise RuntimeError, msg
492
493class PDBReader:
494    """
495    PDB reader class: limited for reading the lines starting with 'ATOM'
496    """
497    type_name = "PDB"
498    ## Wildcards
499    type = ["pdb files (*.PDB, *.pdb)|*.pdb"]
500    ## List of allowed extensions
501    ext = ['.pdb', '.PDB']   
502   
503    def read(self, path):
504        """
505        Load data file
506       
507        :param path: file path
508       
509        :return: MagSLD
510       
511        :raise RuntimeError: when the file can't be opened
512        """
513        pos_x = numpy.zeros(0)
514        pos_y = numpy.zeros(0)
515        pos_z = numpy.zeros(0)
516        sld_n = numpy.zeros(0)
517        sld_mx = numpy.zeros(0)
518        sld_my = numpy.zeros(0)
519        sld_mz = numpy.zeros(0)
520        vol_pix = numpy.zeros(0)
521        pix_symbol = numpy.zeros(0)
522        x_line = []
523        y_line = []
524        z_line = []
525        x_lines = []
526        y_lines = []
527        z_lines = []       
528        try:
529            input_f = open(path, 'rb')
530            buff = input_f.read()
531            lines = buff.split('\n')
532            input_f.close()
533            pre_num = 0
534            num = 0
535            for line in lines:
536                try:
537                    # check if line starts with "ATOM"
538                    if line[0:6].strip().count('ATM') > 0 or \
539                                line[0:6].strip() == 'ATOM':
540                        # define fields of interest
541                        atom_id = line[6:11].strip()
542                        atom_name = line[12:16].strip()
543                        try:
544                            float(line[12])
545                            atom_name = atom_name[1].upper()
546                        except:
547                            if len(atom_name) == 4:
548                                atom_name = atom_name[0].upper()
549                            elif line[12] != ' ': 
550                                atom_name = atom_name[0].upper() + \
551                                        atom_name[1].lower() 
552                            else:
553                                atom_name = atom_name[0].upper()
554                        #res_name = line[17:20].strip()
555                        #chain_name = line[21:22].strip()
556                        _pos_x = float(line[30:38].strip())
557                        _pos_y = float(line[38:46].strip())
558                        _pos_z = float(line[46:54].strip())
559                       
560                        pos_x = numpy.append(pos_x, _pos_x)
561                        pos_y = numpy.append(pos_y, _pos_y)
562                        pos_z = numpy.append(pos_z, _pos_z)
563                        try:
564                            # sld in Ang unit (from fm)
565                            #val = formula(atom_name).atoms.keys()[0].neutron.b_c
566                            #val *= 1.0e-5
567                            val = nsf.neutron_sld(atom_name)[0]
568                            # sld in Ang^-2 unit
569                            val *= 1.0e-6
570                            sld_n = numpy.append(sld_n, val)
571                            atom = formula(atom_name)
572                            # cm to A units
573                            vol = 1.0e+24 * atom.mass / atom.density / NA
574                            vol_pix = numpy.append(vol_pix, vol)
575                        except:
576                            print "Error: set the sld of %s to zero"% atom_name
577                            sld_n = numpy.append(sld_n,  0.0)                     
578                        sld_mx = numpy.append(sld_mx, 0)
579                        sld_my = numpy.append(sld_my, 0)
580                        sld_mz = numpy.append(sld_mz, 0)
581                        pix_symbol = numpy.append(pix_symbol, atom_name)
582                    elif line[0:6].strip().count('CONECT') > 0 :
583                        toks = line.split()
584                        num = int(toks[1]) - 1 
585                        val_list = [] 
586                        for val in toks[2:]:
587                            try:
588                                int_val = int(val)
589                            except:
590                                break
591                            if int_val == 0:
592                                break
593                            val_list.append(int_val)
594                        #need val_list ordered     
595                        for val in val_list:
596                            index = val - 1
597                            if (pos_x[index], pos_x[num]) in x_line and \
598                               (pos_y[index], pos_y[num]) in y_line and \
599                               (pos_z[index], pos_z[num]) in z_line:
600                                continue
601                            x_line.append((pos_x[num], pos_x[index]))
602                            y_line.append((pos_y[num], pos_y[index]))
603                            z_line.append((pos_z[num], pos_z[index]))
604                    if len(x_line) > 0:
605                        x_lines.append(x_line) 
606                        y_lines.append(y_line)
607                        z_lines.append(z_line)           
608                except:
609                    pass
610            #re-centering
611            #pos_x -= (min(pos_x) + max(pos_x)) / 2.0
612            #pos_y -= (min(pos_y) + max(pos_y)) / 2.0
613            #pos_z -= (min(pos_z) + max(pos_z)) / 2.0
614           
615            output = MagSLD(pos_x, pos_y, pos_z, sld_n, sld_mx, sld_my, sld_mz)
616            output.set_conect_lines(x_line, y_line, z_line)
617            output.filename = os.path.basename(path)
618            output.set_pix_type('atom')
619            output.set_pixel_symbols(pix_symbol)
620            output.set_nodes()
621            output.set_pixel_volumes(vol_pix)
622            output.sld_unit = '1/A^(2)'
623            return output
624        except:
625            RuntimeError, "%s is not a sld file" % path
626
627    def write(self, path, data):
628        """
629        Write
630        """
631        #Not implemented
632        print "Not implemented... "
633
634class SLDReader:
635    """
636    Class to load ascii files (7 columns).
637    """
638    ## File type
639    type_name = "SLD ASCII"
640   
641    ## Wildcards
642    type = ["sld files (*.SLD, *.sld)|*.sld",
643            "txt files (*.TXT, *.txt)|*.txt",
644            "all files (*.*)|*.*"]
645    ## List of allowed extensions
646    ext = ['.sld', '.SLD', '.txt', '.TXT', '.*']
647   
648    def read(self, path):
649        """
650        Load data file
651       
652        :param path: file path
653       
654        :return MagSLD: x, y, z, sld_n, sld_mx, sld_my, sld_mz
655       
656        :raise RuntimeError: when the file can't be opened
657        :raise ValueError: when the length of the data vectors are inconsistent
658        """
659        try:
660            pos_x = numpy.zeros(0)
661            pos_y = numpy.zeros(0)
662            pos_z = numpy.zeros(0)
663            sld_n = numpy.zeros(0)
664            sld_mx = numpy.zeros(0)
665            sld_my = numpy.zeros(0)
666            sld_mz = numpy.zeros(0)
667            try:
668                # Use numpy to speed up loading
669                input_f = numpy.loadtxt(path, dtype='float', skiprows=1, 
670                                        ndmin=1, unpack=True)
671                pos_x = numpy.array(input_f[0])
672                pos_y = numpy.array(input_f[1])
673                pos_z = numpy.array(input_f[2])
674                sld_n = numpy.array(input_f[3])
675                sld_mx = numpy.array(input_f[4])
676                sld_my = numpy.array(input_f[5])
677                sld_mz = numpy.array(input_f[6])
678                ncols = len(input_f)
679                if ncols == 8:
680                    vol_pix = numpy.array(input_f[7])
681                elif ncols == 7:
682                     vol_pix = None
683            except:
684                # For older version of numpy
685                input_f = open(path, 'rb')
686                buff = input_f.read()
687                lines = buff.split('\n')
688                input_f.close()
689                for line in lines:
690                    toks = line.split()
691                    try:
692                        _pos_x = float(toks[0])
693                        _pos_y = float(toks[1])
694                        _pos_z = float(toks[2])
695                        _sld_n = float(toks[3])
696                        _sld_mx = float(toks[4])
697                        _sld_my = float(toks[5])
698                        _sld_mz = float(toks[6])
699   
700                        pos_x = numpy.append(pos_x, _pos_x)
701                        pos_y = numpy.append(pos_y, _pos_y)
702                        pos_z = numpy.append(pos_z, _pos_z)
703                        sld_n = numpy.append(sld_n, _sld_n)
704                        sld_mx = numpy.append(sld_mx, _sld_mx)
705                        sld_my = numpy.append(sld_my, _sld_my)
706                        sld_mz = numpy.append(sld_mz, _sld_mz)
707                        try:
708                            _vol_pix =  float(toks[7])
709                            vol_pix = numpy.append(vol_pix, _vol_pix)
710                        except:
711                            vol_pix = None
712                    except:
713                        # Skip non-data lines
714                        pass
715           
716            # re-center     
717            #pos_x -= (min(pos_x) + max(pos_x)) / 2.0
718            #pos_y -= (min(pos_y) + max(pos_y)) / 2.0
719            #pos_z -= (min(pos_z) + max(pos_z)) / 2.0
720
721            output = MagSLD(pos_x, pos_y, pos_z, sld_n, 
722                            sld_mx, sld_my, sld_mz)
723
724            output.filename = os.path.basename(path)
725            output.set_pix_type('pixel')
726            output.set_pixel_symbols('pixel')
727            if vol_pix != None:
728                output.set_pixel_volumes(vol_pix)
729            return output
730        except:
731            RuntimeError, "%s is not a sld file" % path
732       
733    def write(self, path, data):
734        """
735        Write sld file
736       
737        :Param path: file path
738        :Param data: MagSLD data object
739        """
740        if path == None:
741            raise ValueError, "Missing the file path."
742        if data == None:
743            raise ValueError, "Missing the data to save."
744       
745        x_val = data.pos_x
746        y_val = data.pos_y
747        z_val = data.pos_z
748        vol_pix = data.vol_pix
749       
750        length = len(x_val)
751       
752        sld_n = data.sld_n
753        if sld_n == None:
754            sld_n = numpy.zerros(length)
755           
756        sld_mx = data.sld_mx
757        if sld_mx == None:
758            sld_mx = numpy.zerros(length)
759            sld_my = numpy.zerros(length)
760            sld_mz = numpy.zerros(length)
761        else:
762            sld_my = data.sld_my
763            sld_mz = data.sld_mz
764           
765        out = open(path, 'w')
766        # First Line: Column names
767        out.write("X  Y  Z  SLDN SLDMx  SLDMy  SLDMz VOLUMEpix")
768        for ind in range(length):
769            out.write("\n%g  %g  %g  %g  %g  %g  %g %g" % (x_val[ind], y_val[ind], 
770                                                    z_val[ind], sld_n[ind], 
771                                                    sld_mx[ind], sld_my[ind], 
772                                                    sld_mz[ind], vol_pix[ind])) 
773        out.close()       
774           
775           
776class OMFData:
777    """
778    OMF Data.
779    """
780    _meshunit = "A"
781    _valueunit = "A^(-2)"
782    def __init__(self):
783        """
784        Init for mag SLD
785        """
786        self.filename = 'default'
787        self.oommf = ''
788        self.title = ''
789        self.desc = ''
790        self.meshtype = ''
791        self.meshunit = self._meshunit
792        self.valueunit = self._valueunit
793        self.xbase = 0.0
794        self.ybase = 0.0
795        self.zbase = 0.0
796        self.xstepsize = 6.0
797        self.ystepsize = 6.0
798        self.zstepsize = 6.0
799        self.xnodes = 10.0
800        self.ynodes = 10.0
801        self.znodes = 10.0
802        self.xmin = 0.0
803        self.ymin = 0.0
804        self.zmin = 0.0
805        self.xmax = 60.0
806        self.ymax = 60.0
807        self.zmax = 60.0
808        self.mx = None
809        self.my = None
810        self.mz = None
811        self.valuemultiplier = 1.
812        self.valuerangeminmag = 0
813        self.valuerangemaxmag = 0
814       
815    def __str__(self):
816        """
817        doc strings
818        """
819        _str =  "Type:            %s\n" % self.__class__.__name__
820        _str += "File:            %s\n" % self.filename
821        _str += "OOMMF:           %s\n" % self.oommf
822        _str += "Title:           %s\n" % self.title
823        _str += "Desc:            %s\n" % self.desc
824        _str += "meshtype:        %s\n" % self.meshtype
825        _str += "meshunit:        %s\n" % str(self.meshunit)
826        _str += "xbase:           %s [%s]\n" % (str(self.xbase), self.meshunit)
827        _str += "ybase:           %s [%s]\n" % (str(self.ybase), self.meshunit)
828        _str += "zbase:           %s [%s]\n" % (str(self.zbase), self.meshunit)
829        _str += "xstepsize:       %s [%s]\n" % (str(self.xstepsize), 
830                                                self.meshunit)
831        _str += "ystepsize:       %s [%s]\n" % (str(self.ystepsize), 
832                                                self.meshunit)
833        _str += "zstepsize:       %s [%s]\n" % (str(self.zstepsize), 
834                                                self.meshunit)
835        _str += "xnodes:          %s\n" % str(self.xnodes)
836        _str += "ynodes:          %s\n" % str(self.ynodes)
837        _str += "znodes:          %s\n" % str(self.znodes)
838        _str += "xmin:            %s [%s]\n" % (str(self.xmin), self.meshunit)
839        _str += "ymin:            %s [%s]\n" % (str(self.ymin), self.meshunit)
840        _str += "zmin:            %s [%s]\n" % (str(self.zmin), self.meshunit)
841        _str += "xmax:            %s [%s]\n" % (str(self.xmax), self.meshunit)
842        _str += "ymax:            %s [%s]\n" % (str(self.ymax), self.meshunit)
843        _str += "zmax:            %s [%s]\n" % (str(self.zmax), self.meshunit)
844        _str += "valueunit:       %s\n" % self.valueunit
845        _str += "valuemultiplier: %s\n" % str(self.valuemultiplier)
846        _str += "ValueRangeMinMag:%s [%s]\n" % (str(self.valuerangeminmag), 
847                                                 self.valueunit)
848        _str += "ValueRangeMaxMag:%s [%s]\n" % (str(self.valuerangemaxmag), 
849                                                 self.valueunit)
850        return _str
851   
852    def set_m(self, mx, my, mz):
853        """
854        Set the Mx, My, Mz values
855        """
856        self.mx = mx
857        self.my = my
858        self.mz = mz
859               
860class MagSLD:
861    """
862    Magnetic SLD.
863    """
864    pos_x = None
865    pos_y = None
866    pos_z = None
867    sld_n = None
868    sld_mx = None
869    sld_my = None
870    sld_mz = None
871    # Units
872    _pos_unit = 'A'
873    _sld_unit = '1/A^(2)'
874    _pix_type = 'pixel'
875   
876    def __init__(self, pos_x, pos_y, pos_z, sld_n=None, 
877                 sld_mx=None, sld_my=None, sld_mz=None, vol_pix=None):
878        """
879        Init for mag SLD
880        :params : All should be numpy 1D array
881        """
882        self.is_data = True
883        self.filename = ''
884        self.xstepsize = 6.0
885        self.ystepsize = 6.0
886        self.zstepsize = 6.0
887        self.xnodes = 10.0
888        self.ynodes = 10.0
889        self.znodes = 10.0
890        self.has_stepsize = False
891        self.has_conect = False
892        self.pos_unit = self._pos_unit
893        self.sld_unit = self._sld_unit
894        self.pix_type = 'pixel'
895        self.pos_x = pos_x
896        self.pos_y = pos_y
897        self.pos_z = pos_z
898        self.sld_n = sld_n
899        self.line_x = None
900        self.line_y = None
901        self.line_z = None
902        self.sld_mx = sld_mx
903        self.sld_my = sld_my
904        self.sld_mz = sld_mz
905        self.vol_pix = vol_pix
906        self.sld_m = None
907        self.sld_phi = None
908        self.sld_theta = None
909        self.pix_symbol = None
910        if sld_mx != None and sld_my != None and sld_mz != None:
911            self.set_sldms(sld_mx, sld_my, sld_mz)
912        self.set_nodes()
913                   
914    def __str__(self):
915        """
916        doc strings
917        """
918        _str =  "Type:       %s\n" % self.__class__.__name__
919        _str += "File:       %s\n" % self.filename
920        _str += "Axis_unit:  %s\n" % self.pos_unit
921        _str += "SLD_unit:   %s\n" % self.sld_unit
922        return _str
923   
924    def set_pix_type(self, pix_type):
925        """
926        Set pixel type
927        :Param pix_type: string, 'pixel' or 'atom'
928        """
929        self.pix_type = pix_type
930       
931    def set_sldn(self, sld_n):
932        """
933        Sets neutron SLD
934        """
935        if sld_n.__class__.__name__ == 'float':
936            if self.is_data:
937                # For data, put the value to only the pixels w non-zero M
938                is_nonzero = (numpy.fabs(self.sld_mx) + 
939                                  numpy.fabs(self.sld_my) + 
940                                  numpy.fabs(self.sld_mz)).nonzero() 
941                self.sld_n = numpy.zeros(len(self.pos_x))
942                if len(self.sld_n[is_nonzero]) > 0:
943                    self.sld_n[is_nonzero] = sld_n
944                else:
945                    self.sld_n.fill(sld_n)
946            else:
947                # For non-data, put the value to all the pixels
948                self.sld_n = numpy.ones(len(self.pos_x)) * sld_n
949        else:
950            self.sld_n = sld_n
951           
952    def set_sldms(self, sld_mx, sld_my, sld_mz):
953        """
954        Sets (|m|, m_theta, m_phi)
955        """
956        if sld_mx.__class__.__name__ == 'float':
957            self.sld_mx = numpy.ones(len(self.pos_x)) * sld_mx
958        else:
959            self.sld_mx = sld_mx
960        if sld_my.__class__.__name__ == 'float':
961            self.sld_my = numpy.ones(len(self.pos_x)) * sld_my
962        else:
963            self.sld_my = sld_my
964        if sld_mz.__class__.__name__ == 'float':
965            self.sld_mz = numpy.ones(len(self.pos_x)) * sld_mz
966        else:
967            self.sld_mz = sld_mz
968
969        sld_m = numpy.sqrt(sld_mx * sld_mx + sld_my * sld_my + \
970                                sld_mz * sld_mz) 
971        self.sld_m = sld_m
972   
973    def set_pixel_symbols(self, symbol='pixel'): 
974        """
975        Set pixel
976        :Params pixel: str; pixel or atomic symbol, or array of strings
977        """
978        if self.sld_n == None:
979            return
980        if symbol.__class__.__name__ == 'str':
981            self.pix_symbol = numpy.repeat(symbol, len(self.sld_n))
982        else:
983            self.pix_symbol = symbol
984           
985    def set_pixel_volumes(self, vol): 
986        """
987        Set pixel volumes
988        :Params pixel: str; pixel or atomic symbol, or array of strings
989        """
990        if self.sld_n == None:
991            return
992        if vol.__class__.__name__ == 'ndarray':
993            self.vol_pix = vol
994        elif vol.__class__.__name__.count('float') > 0:
995            self.vol_pix = numpy.repeat(vol, len(self.sld_n))
996        else:
997            self.vol_pix = None
998
999    def get_sldn(self):
1000        """
1001        Returns nuclear sld
1002        """
1003        return self.sld_n
1004   
1005    def get_sld_mxyz(self):
1006        """
1007        Returns (sld_mx, sldmy, sld_mz)
1008        """
1009        return (self.sld_mx, self.sldmy, self.sld_mz)
1010   
1011    def get_sld_m(self):
1012        """
1013        Returns (sldm, sld_theta, sld_phi)
1014        """
1015        return (self.sldm, self.sld_theta, self.sld_phi)
1016   
1017    def set_nodes(self):
1018        """
1019        Set xnodes, ynodes, and znodes
1020        """
1021        self.set_stepsize()
1022        if self.pix_type == 'pixel':
1023            try:
1024                xdist = (max(self.pos_x) - min(self.pos_x)) / self.xstepsize
1025                ydist = (max(self.pos_y) - min(self.pos_y)) / self.ystepsize
1026                zdist = (max(self.pos_z) - min(self.pos_z)) / self.zstepsize
1027                self.xnodes = int(xdist) + 1
1028                self.ynodes = int(ydist) + 1
1029                self.znodes = int(zdist) + 1
1030            except:
1031                self.xnodes = None
1032                self.ynodes = None
1033                self.znodes = None
1034        else:
1035            self.xnodes = None
1036            self.ynodes = None
1037            self.znodes = None
1038     
1039    def set_stepsize(self):
1040        """
1041        Set xtepsize, ystepsize, and zstepsize
1042        """
1043        if self.pix_type == 'pixel':
1044            try:
1045                xpos_pre = self.pos_x[0]
1046                ypos_pre = self.pos_y[0]
1047                zpos_pre = self.pos_z[0]
1048                for x_pos in self.pos_x:
1049                    if xpos_pre != x_pos:
1050                        self.xstepsize = numpy.fabs(x_pos - xpos_pre)
1051                        break
1052                for y_pos in self.pos_y:
1053                    if ypos_pre != y_pos:
1054                        self.ystepsize = numpy.fabs(y_pos - ypos_pre)
1055                        break
1056                for z_pos in self.pos_z:
1057                    if zpos_pre != z_pos:
1058                        self.zstepsize = numpy.fabs(z_pos - zpos_pre)
1059                        break
1060                #default pix volume
1061                self.vol_pix = numpy.ones(len(self.pos_x))
1062                vol = self.xstepsize * self.ystepsize * self.zstepsize
1063                self.set_pixel_volumes(vol)
1064                self.has_stepsize = True
1065            except:
1066                self.xstepsize = None
1067                self.ystepsize = None
1068                self.zstepsize = None
1069                self.vol_pix = None
1070                self.has_stepsize = False
1071        else:
1072            self.xstepsize = None
1073            self.ystepsize = None
1074            self.zstepsize = None
1075            self.has_stepsize = True
1076        return self.xstepsize, self.ystepsize, self.zstepsize
1077
1078    def set_conect_lines(self, line_x, line_y, line_z):
1079        """
1080        Set bonding line data if taken from pdb
1081        """
1082        if line_x.__class__.__name__ != 'list' or len(line_x) < 1:
1083            return
1084        if line_y.__class__.__name__ != 'list' or len(line_y) < 1:
1085            return
1086        if line_z.__class__.__name__ != 'list' or len(line_z) < 1:
1087            return
1088        self.has_conect = True
1089        self.line_x = line_x
1090        self.line_y = line_y
1091        self.line_z = line_z
1092       
1093
1094def test_load():
1095    from danse.common.plottools.arrow3d import Arrow3D
1096    import os
1097    dir = os.path.abspath(os.path.curdir)
1098    print dir
1099    for i in range(6):
1100        dir, _ = os.path.split(dir)
1101        tfile = os.path.join(dir, "test", "CoreXY_ShellZ.txt")
1102        ofile = os.path.join(dir, "test", "A_Raw_Example-1.omf")
1103        if os.path.isfile(tfile):
1104            tfpath = tfile
1105            ofpath = ofile
1106            break
1107    reader = SLDReader()
1108    oreader = OMFReader()
1109    output = reader.read(tfpath)
1110    ooutput = oreader.read(ofpath)
1111    foutput = OMF2SLD()
1112    foutput.set_data(ooutput)
1113
1114    import matplotlib.pyplot as plt
1115    from mpl_toolkits.mplot3d import Axes3D
1116    fig = plt.figure()
1117    ax = fig.gca(projection='3d')
1118    ax.plot(output.pos_x, output.pos_y, output.pos_z, '.', c="g", 
1119            alpha = 0.7, markeredgecolor='gray',rasterized=True)
1120    gap = 7
1121    max_mx = max(output.sld_mx)
1122    max_my = max(output.sld_my)
1123    max_mz = max(output.sld_mz)
1124    max_m = max(max_mx,max_my,max_mz)
1125    x2 = output.pos_x+output.sld_mx/max_m * gap
1126    y2 = output.pos_y+output.sld_my/max_m * gap
1127    z2 = output.pos_z+output.sld_mz/max_m * gap
1128    x_arrow = numpy.column_stack((output.pos_x,x2))
1129    y_arrow = numpy.column_stack((output.pos_y,y2))
1130    z_arrow = numpy.column_stack((output.pos_z,z2))
1131    unit_x2 = output.sld_mx / max_m
1132    unit_y2 = output.sld_my / max_m
1133    unit_z2 = output.sld_mz / max_m
1134    color_x = numpy.fabs(unit_x2 * 0.8)
1135    color_y = numpy.fabs(unit_y2 * 0.8)
1136    color_z = numpy.fabs(unit_z2 * 0.8)
1137    colors =  numpy.column_stack((color_x, color_y, color_z))
1138    a = Arrow3D(None, x_arrow,y_arrow,z_arrow, colors, mutation_scale=10, lw=1, 
1139                arrowstyle="->", color="y", alpha = 0.5)
1140
1141    ax.add_artist(a)
1142    plt.show()
1143   
1144def test():
1145    import os
1146    dir = os.path.abspath(os.path.curdir)
1147    for i in range(3):
1148        dir, _ = os.path.split(dir)
1149        #tfile = os.path.join(dir, "test", "C_Example_Converted.txt")
1150        ofile = os.path.join(dir, "test", "A_Raw_Example-1.omf")
1151        if os.path.isfile(ofile):
1152            #tfpath = tfile
1153            ofpath = ofile
1154            break
1155    #reader = SLDReader()
1156    oreader = OMFReader()
1157    #output = reader.read(tfpath)
1158    ooutput = oreader.read(ofpath)
1159    foutput = OMF2SLD() 
1160    foutput.set_data(ooutput) 
1161    writer =  SLDReader() 
1162    writer.write(os.path.join(os.path.dirname(ofpath), "out.txt"), 
1163                 foutput.output)
1164    model = GenSAS()
1165    model.set_sld_data(foutput.output) 
1166    x = numpy.arange(1000)/10000. + 1e-5
1167    y = numpy.arange(1000)/10000. + 1e-5
1168    #z = numpy.arange(1000)/10000. + 1e-5
1169    i = numpy.zeros(1000)
1170    out = model.runXY([x,y,i])
1171       
1172if __name__ == "__main__": 
1173    test()
1174    test_load()
Note: See TracBrowser for help on using the repository browser.