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

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 fdb134b9 was fdb134b9, checked in by Jae Cho <jhjcho@…>, 11 years ago

added bonding lines for pdb data if provided

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