import os,sys,copy from xml.etree import ElementTree as ET import inout,common class files(object): # mtz labels handling should be moved here (?) def __init__(self,fname,ftyp): self.name=fname self.typ=ftyp def Data2XML(self, ET_element): """Convert the actual file object into XML""" xroot = ET.SubElement(ET_element, 'file') xroot.text = self.name xroot.attrib['typ']=self.typ def Convert(self,newtyp,par_proc,inp_cont): self.no_conversion = False if all(c.GetFile(newtyp) for c in inp_cont): return inp_cont[0].GetFile(newtyp) from_to=(self.typ,newtyp) # in case we are called from program, find its parent process try: par_proc=par_proc.process except AttributeError: sys.exc_clear() # create the convert process and run it convert_proc=par_proc.AddProcess('convert',propagate_inp=False,propagate_out=False) if from_to not in convert_proc.setup: par_proc.processes.remove(convert_proc) return None for cont in inp_cont: convert_proc.inp.Add(cont) convert_proc.SetRunDir() converted_file=None if not self.no_conversion: try: convert_proc.setup[from_to](convert_proc,*from_to) convert_proc.Run(clear_out=False) except Exception as e: #common.Warning('Conversion from {0} to {1} of process {2} failed with message: {3}'.format(self.typ,newtyp,par_proc.nick,e.message)) sys.exc_clear() else: converted_file = convert_proc.out.Get(filetype=newtyp).GetFile(newtyp) # add the created file to the container parent if needed #for cont in inp_cont: # if not cont.GetFile(newtyp): # cont.AddFile( converted_file.name, newtyp ) # clean up and return par_proc.processes.remove(convert_proc) return converted_file class data_container(object): """base container for data (observed and calculated) storage classes""" # the internal booking changes - cell, spacegroup etc might be treated as property of crystal etc # mtz column labels should be hidden on the files level etc default_unknown='Undefined' deffiletype=None supported_filetypes=[] xname=default_unknown dname=default_unknown cell=None spgr=None spgr_num=None resol=None col_list=[] label_default=None labelout_prefix='' custom=[] typ='' solvent_content=None monomers_asym=None residues_mon=None # list of supported arguments set by user (as opposed to set by crank) # this might get tricky when new objects are created from objects with user inputted args... set_by_user = [] def __init__(self, filename=None, filetype=None, typ=None, lab_dict={}, custom=None, \ xname=None, dname=None, xmlelem=None, inpline=None): self.nick=self.__class__.__name__ self.mtzfile_asked=False if typ is not None: self.typ=typ self.files=[] if filename is not None: self.AddFile(filename, filetype) if self.col_list: self.InitLabels(lab_dict) self.custom=[] if custom: self.custom.append(custom) if xname: self.xname=xname if dname: self.dname=dname self.delayed_init={} if xmlelem is not None: self.XMLElemInit(xmlelem) elif inpline is not None: self.InputElemInit(inpline) @classmethod def from_name(cls, container_name, xmlelem=None, inpline=None): """Create specific (derived) container instance from the (nick)name of the container""" try: inst = getattr(sys.modules[__name__], container_name)(xmlelem=xmlelem,inpline=inpline) except (AttributeError,KeyError): common.Error('Error when creating container {0} instance.'.format(container_name)) else: return inst def DelayedInit(self, inp, numstep, crank): """Performs delayed initialization of this object: replaces this object by the specified previously outputted object. Leaves gracefully if this object does not contain any delayed init. information. It is an error if this object contains delayed init. information but the delayed init. failed. Arguments: inp - the program/process input asking this object's delayed initialization numstep - order number of the step asking this object's delayed initialization crank - the main crank process """ if self.delayed_init: step=self.delayed_init.pop('step') if step>=numstep: common.Error('Delayed init of {0} in step {1} failed: requested step {2} has not happened yet'.format( self.nick, numstep, step)) if not crank.run.processes[step].out.Get(**self.delayed_init): common.Error('Delayed initialization of {0} in step {1} failed: no such object from step {2}.'.format( self.nick, numstep, step)) inp.AddCopy(crank.run.processes[step].out.Get(**self.delayed_init), propagate=False) self.delayed_init={} inp.Delete(self) def XMLElemInit(self,xelem): """Initialize from inputted XML element""" for atr,val in xelem.attrib.iteritems(): try: getattr(self,atr) except AttributeError: common.Error('Unknown attribute {0} for data container {1}'.format(atr,self.nick)) else: setattr(self,atr,val) attr_set_by_crank=[] tags = list(self.supported_attributes) + [ "delayed_init","attr_set_by_crank" ] + self.col_list for xchild in xelem: if xchild.tag in tags: if xchild.tag=='attr_set_by_crank': attr_set_by_crank = common.AutoConvert(xchild.text) elif xchild.tag=='file': filetyp=None if 'typ' in xchild.attrib: filetyp=xchild.attrib['typ'] self.AddFile(xchild.text,filetyp) elif xchild.tag=='atomtype': fp,fpp=None,None for xatomchild in xchild: if xatomchild.tag=='fp': fp=common.AutoConvert(xatomchild.text) elif xatomchild.tag=='fpp': fpp=common.AutoConvert(xatomchild.text) else: common.Error('Wrong subtag {0} of atomtype when initiating data container {1} from XML'.format( atomxchild.tag,self.nick)) self.SetAtomType(xchild.text.strip(),fp,fpp) elif xchild.tag in ('xname','dname'): setattr(self,xchild.tag,xchild.text) elif xchild.tag in ('solvent_content','monomers_asym','residues_mon','exp_num_atoms','custom'): setattr(self,xchild.tag,common.AutoConvert(xchild.text)) elif xchild.tag=='delayed_init': for xdelinichild in xchild: self.delayed_init[xdelinichild.tag] = common.AutoConvert(xdelinichild.text) elif xchild.tag in self.col_list: self.SetLabel(xchild.tag,xchild.text) if xchild.tag in self.supported_attributes and xchild.tag not in attr_set_by_crank and \ xchild.tag not in self.set_by_user: self.set_by_user.append(xchild.tag) else: common.Error('Wrong tag {0} when initializing data container {1} from XML'.format(xchild.tag,self.nick)) if 'typ' not in attr_set_by_crank and 'typ' not in self.set_by_user: self.set_by_user.append('typ') def InputElemInit(self,input_line): """Initialize from input line""" ad={'atomtype':None,'fp':None,'fpp':None} keys = list(self.supported_attributes) + [ "obj_from", ] + self.col_list while input_line: token=input_line.pop(0) if token in self.types: self.SetType(token) if not 'typ' in self.set_by_user: self.set_by_user.append('typ') else: key,sep,val = token.partition('=') if key not in keys: input_line.insert(0,token) break if key=='typ': self.SetType(val) elif key in ('xname','dname'): setattr(self,key,val) elif key in ('solvent_content','monomers_asym','residues_mon','exp_num_atoms','custom'): setattr(self,key,common.AutoConvert(val)) # filetype is always determined from suffix as of now! elif key=='file': filetyp=val.split('.')[-1] if not os.path.isabs(val): val = os.path.join(os.getcwd(), val) self.AddFile(val,filetyp) elif key in ('atomtype','fp','fpp'): ad[key]=val if ad['atomtype']: self.SetAtomType(ad['atomtype'],ad['fp'],ad['fpp']) if ad['fp'] and ad['fpp']: ad={'atomtype':None,'fp':None,'fpp':None} elif key in self.col_list: self.SetLabel(key,val) elif key=='obj_from': args = val.split(',') if not args or type(common.AutoConvert(args[0]))!=int: common.Error("Input keyword {0} needs an integer specifying the step".format(key)) self.delayed_init['step']=common.AutoConvert(args[0]) for cond in args[1:]: k,s,v=cond.partition('=') if not v: common.Error("Condition {0} of input keyword {1} malformed.".format(cond,key)) if k not in getattr(inout.input_output,'GetAll').func_code.co_varnames: common.Error("Wrong subkeyword {0} of input keyword {1}.".format(k,key)) self.delayed_init[k]=common.AutoConvert(v) if not key in self.set_by_user: self.set_by_user.append(key) def Data2XML(self, ET_element): """Convert the actual container object into XML""" xroot = ET.SubElement(ET_element, self.__class__.__name__) xroot.attrib['typ']=self.typ attr_set_by_crank = [] for attr in self.supported_attributes: if hasattr(self,attr) and getattr(self,attr) and attr not in self.set_by_user: attr_set_by_crank.append(attr) if attr_set_by_crank: ET.SubElement(xroot,"attr_set_by_crank").text = str(attr_set_by_crank) if self.dname != self.default_unknown: ET.SubElement(xroot,'dname').text = self.dname if self.xname != self.default_unknown: ET.SubElement(xroot,'xname').text = self.xname for attr in ('solvent_content','monomers_asym','residues_mon','custom','exp_num_atoms'): if hasattr(self,attr) and getattr(self,attr): ET.SubElement(xroot,attr).text = str(getattr(self,attr)) for f in self.files: f.Data2XML(xroot) for col,lab in self.GetAllLabels(): ET.SubElement(xroot,col).text = lab if 'atomtypes' in self.__dict__: atomtypes = [self.atomtype1,] + [ at for at in self.atomtypes if at!=self.atomtype1 ] for at in [at for at in atomtypes if at is not None]: xat = ET.SubElement(xroot,'atomtype') xat.text = at if self.atomtypes[at]: ET.SubElement(xat,'fp').text=str(self.atomtypes[at][0]) ET.SubElement(xat,'fpp').text=str(self.atomtypes[at][1]) if self.delayed_init: xdelini = ET.SubElement(xroot,'delayed_init') for k,v in self.delayed_init.iteritems(): ET.SubElement(xdelini,k).text=str(v) def IsInputtedAttr(self, attr): """Returns True if the data object attribute has been inputted by user""" if attr in self.set_by_user: return True return False def GetFiles(self, filetype=None, conv_parent=None, inp_cont=None): """Returns list of file objects of the specified file type associated with this container Arguments: filetype - type of the files to be returned; if None or omitted then all files are returned; a vector/list of string filetypes is also accepted meaning logical OR conv_parent - specifies that conversion of file should be tried and parent process or program that asks it (in case of multiple filetypes, only the first is attempted for conversion) inp_cont - list of input containers (other than this container) for conversion (single container object also accepted) """ if isinstance(filetype,basestring): filetype=[filetype,] result = [ f for f in self.files if filetype is None or f.typ in filetype ] if not result and conv_parent: if inp_cont is None: inp_cont=[] try: inp_cont.append(self) except AttributeError: inp_cont=[inp_cont,self] sys.exc_clear() for f in self.files: if f.Convert(filetype[0],conv_parent,inp_cont): break result = [ f for f in self.files if f.typ==filetype[0] ] return result def GetFile(self, filetype=None, conv_parent=None, inp_cont=None, ind=-1): """Returns container's file object with the specified file type; None if there is none; in case of multiple files of the file type, 'ind' specifies the index (last by default) The same arguments as GetFiles() + the ind parameter """ if inp_cont is None: inp_cont=[] try: return self.GetFiles(filetype,conv_parent,inp_cont)[ind] except IndexError: sys.exc_clear() return None def GetFileName(self, filetype=None, conv_parent=None, inp_cont=None, ind=-1, trim_path=False): """Returns filename of a file object of the specified file type; None if there is none In case of multiple files of the file type, 'ind' specifies the index (last by default) The same arguments as GetFiles() + the ind parameter """ try: f=self.GetFiles(filetype, conv_parent, inp_cont)[ind] except IndexError: sys.exc_clear() return None else: if f.typ=='mtz': self.mtzfile_asked=True fname = f.name if trim_path: fname = os.path.basename(fname) return fname def SetLabel(self, col, lbl=None, ignore_prefix=False): """Set (mtz) column(s) to label(s) Arguments: col - column(s) whose labels should be set; can be either a string or a list of strings lbl - label(s) to be set, needs to be of the same type as col (and the same order assumed for list) if lbl is None or omitted then the default label is set (if defined, otherwise it is an error) """ # if string inputted then convert to list if isinstance(col,basestring): if lbl and not isinstance(lbl,basestring): common.Error('Wrong argument types when setting label for {0}'.format(self.nick)) col=[col,] if lbl: lbl=[lbl,] # check the inputted col for c in col: if c not in self.col_list: common.Error('No column list or no column {0} defined for {1}.'.format(c,self.nick)) # assign defaults if asked for if not lbl: if not self.label_default or self.typ not in self.label_default: common.Error('No default label definition for {0} of type {1}'.format(self.nick, self.typ)) lbl = [ self.label_default[self.typ][self.col_list.index(c)] for c in col ] # do the actual setting for c,l in zip(col,lbl): if not ignore_prefix: l=self.labelout_prefix+l self.label[c]=l def GetLabel(self, col, mtz_not_asked=False): """Get label of the inputted column""" if not mtz_not_asked: self.mtzfile_asked=True if col not in self.col_list: common.Error('Column {0} not defined for container {1}'.format(col,self.nick)) return self.label[col] def InitLabels(self, lab_dict={}): """Initialize (mtz) label dictionary with keys from internal col_list and (optionally) inputted dictionary values""" self.label={} for col in self.col_list: self.label[col]=None for col in lab_dict: self.SetLabel(col,lab_dict[col]) def GetAllLabels(self, labels_only=False, mtz_not_asked=False): """Returns all defined (non None) labels in a list of pairs (column,label) Returns just a list of (non None) labels if labels_only is True (False by default) """ if not labels_only: return [ (c,self.GetLabel(c,mtz_not_asked)) for c in self.col_list if self.GetLabel(c) is not None ] else: return [ self.GetLabel(c,mtz_not_asked) for c in self.col_list if self.GetLabel(c) is not None ] def GetFullLabel(self, *cols, **kwargs): """Return full (mtz) label path for specified column(s) (ie 'crystal/data/[labels]')""" objs = [self,] + list(kwargs.get('other_objs', [])) if not self.col_list or not self.GetFile('mtz'): common.Error('No column list or no mtz file for {0}.'.format(self.nick)) crystal_name = self.xname if self.xname!=self.default_unknown else '*' data_name = self.dname if self.dname!=self.default_unknown else '*' labels=[] for obj in objs: if obj.xname!=self.xname: common.Warning('xname mismatch when making full label: {0} vs {1}'.format(obj.xname,self.xname)) if obj.dname!=self.dname: common.Warning('dname mismatch when making full label: {0} vs {1}'.format(obj.dname,self.dname)) for col in cols: labels.append(obj.GetLabel(col)) return '/'.join((crystal_name,data_name,'['+','.join(labels)+']')) def SetType(self, typ, preserve_labels=False): """Set type of the data container""" if typ in self.types: self.typ=typ else: common.Error('Type {0} not allowed for container {1} (the following types are allowed: {2})'.format( typ,self.nick,','.join(self.types))) def GetType(self): """Get type of the data container""" return self.typ def AddFile(self, filename, filetype=None): """Add/associate file with filename of filetype to the data container The filename must be supplied as absolute path otherwise there is an error """ assert os.path.isabs(filename), 'File {0} is not specified as absolute path.'.format(filename) if not filetype: filetype=self.deffiletype if filetype not in self.supported_filetypes: common.Error('Wrong filetype {0} for container {1}'.format(filetype,self.nick)) self.files.append(files(filename,filetype)) def SetFile(self, filename, filetype=None): """The same as AddFile() but any files of the same filetype will be removed before adding the file""" try: self.DetachFile(filetype=filetype) except self.NoDetachFileError: sys.exc_clear() self.AddFile(filename,filetype) class NoDetachFileError(Exception): pass def DetachFile(self, filename=None, filetype=None): """Removes association of file with filename and/or file(s) of filetype with the container It is an error if there is no such file """ if not filetype: filetype=self.deffiletype removed=False for f in self.files: if f.typ==filetype or f.name==filename: self.files.remove(f) removed=True if not removed: self.NoDetachFileError('No files with name {0} and/or type {1} could be detached from container {2}'.format( filename, filetype, self.nick) ) def SetCrystalName(self, name): self.xname=name def GetCrystalName(self): return self.xname def SetDataName(self, name): self.dname=name def GetDataName(self): return self.dname def GetCellSpacegroupResol(self, pro, get='', accept_none=False): # pro is here to define the directory (associated to the program/process pro) in which sftools will be run # furthermore, if the information cannot be obtained from this object, it will attempt to get it # from another object from pro.inp with the same xname, dname if possible mtz_obj=None if self.GetFile('mtz'): mtz_obj = self else: mtz_obj = pro.inp.Get(filetype='mtz',xname=self.GetCrystalName(),dname=self.GetDataName()) if not mtz_obj: if accept_none: return None common.Error('Spacegroup or cell or resolution could not be retrieved.') if not mtz_obj.GetLabel('f') and get=='res': mtz_obj = pro.inp.Get(filetype='mtz',xname=self.GetCrystalName(),dname=self.GetDataName(),col='f') if not mtz_obj: if accept_none: return None common.Error('Resolution could not be retrieved.') try: pro=pro.process except AttributeError: sys.exc_clear() while pro.rundir is None: pro=pro.parent_process sft = pro.AddProg('sftools') sft.runname='sftools_getinfo' sft.SetKey('read', self.GetFileName('mtz')) if get=='res': sft.SetKey('select',['col',self.GetLabel('f')]) sft.SetKey('checkhkl') else: sft.SetKey('list') if self.xname!=self.default_unknown: sft.SetKey('list dcell') sft.SetKey('quit\nY') sft.Run() if get=='cell': self.cell = ( sft.GetStat('a',self.xname), sft.GetStat('b',self.xname), \ sft.GetStat('c',self.xname), sft.GetStat('alpha',self.xname), \ sft.GetStat('beta',self.xname), sft.GetStat('gamma',self.xname) ) info = self.cell elif get=='spg': self.spgr=sft.GetStat('spacegroup').strip() info = self.spgr elif get=='spgn': self.spgr_num=sft.GetStat('spacegroup_num') info = self.spgr_num elif get=='res': self.resol=sft.GetStat('resolution',self.GetLabel('f')[:12]) info = self.resol else: common.Error('Internal error in GetCellSpacegroup.') pro.programs.remove(sft) return info def GetCell(self,pro): if self.cell: return self.cell else: return self.GetCellSpacegroupResol(pro,get='cell') def GetSpacegroup(self,pro): if self.spgr: return self.spgr else: return self.GetCellSpacegroupResol(pro,get='spg') def GetSpacegroupNumber(self,pro): if self.spgr_num: return self.spgr_num else: return self.GetCellSpacegroupResol(pro,get='spgn') def GetResolution(self,pro): if self.resol: return self.resol else: return self.GetCellSpacegroupResol(pro,get='res') def AddCustomTag(self,tag): self.custom.append(tag) class fsigf(data_container): """Observed amplitudes/intensities container class""" description="observed amplitudes/intensities" supported_filetypes=('mtz','hkl','sca','drear') deffiletype='mtz' col_list=['f','sigf','i','sigi','alpha'] types=('average','plus','minus','fa','ea','delta-anom') typ='average' label_default = { 'average': ['F','SIGF','I','SIGI',None], \ 'plus': ['F+','SIGF+','I+','SIGI+',None], \ 'minus': ['F-','SIGF-','I-','SIGI-',None], \ 'fa': ['FA','SIGFA',None,None,'ALPHA'], \ 'ea': ['EA','SIGEA',None,None,'ALPHA'], \ 'delta-anom': ['DA','SIGDA','IDA','SIGIDA',None,], \ } supported_attributes=('typ','file','xname','dname','custom') class mapcoef(data_container): """Map coefficients (incl. phase information) container class""" description="map coefficients (incl. phase information)" supported_filetypes=('mtz','map','phs') deffiletype='mtz' col_list=['f', 'ph', 'fom', 'hla', 'hlb', 'hlc', 'hld',] types=('best','weighted','combined','densmod','anomalous','anom-diff') typ='best' label_default = { 'best': ['FB','PHIB','FOM','HLA','HLB','HLC','HLD'], \ 'weighted': ['FWT','PHWT',None,None,None,None,None], \ 'combined': ['FCOMB','PHCOMB','FOMCOMB','HLACOMB','HLBCOMB','HLCCOMB','HLDCOMB'], \ 'densmod': ['FMOD','PHMOD',None,None,None,None,None], \ 'anomalous': ['FAN','PHAN',None,None,None,None,None], \ 'anom-diff': ['DELFAN','PHDELAN',None,None,None,None,None], \ } supported_attributes=('typ','file','xname','dname','custom') class sequence(data_container): """Sequence container class""" description="protein sequence" supported_filetypes=('pir','seq','fasta','fas') deffiletype='pir' types=('protein',) typ='protein' supported_attributes=('file','xname','solvent_content','monomers_asym','residues_mon') # the sequence string, if known. seqstr="" def GetSequenceString(self): """Gets the sequence for the file for internal use/decision making. At the moment, there is no proper support for the different sequence formats implemented, just a simple heuristics is used to grab the sequence. """ self.seqstr="" with open(self.GetFileName(),'r') as f: for il,line in enumerate(f): line=line.strip() if not line.startswith('>') and not line.startswith('#'): self.seqstr+=line; if not line and il==1 and self.GetFile().typ=='pir': self.seqstr="" class model(data_container): """Model (of crystal content) container class""" description="crystal model and its description" # xyz is the PMF substr. output, res is the SHELXD output substr., hat is the SHELXE susbtr. output supported_filetypes=('pdb','xyz','res','hat') deffiletype='pdb' types=('partial','substr','partial+substr','unknown','patterson') typ='unknown' exp_num_atoms=None supported_attributes=('typ','file','xname','atomtype','fp','fpp','solvent_content','monomers_asym','exp_num_atoms','custom') def __init__(self,*args,**kwargs): # dictionary of atom types and their f',f'' self.atomtypes={} self.atomtype1=kwargs.pop('atomtype1',None) for at,f in kwargs.pop('atomtypes',{}).iteritems(): self.SetAtomType(at,f[0],f[1]) super(model,self).__init__(*args,**kwargs) def SetAtomType(self, atomtype, fp=None, fpp=None): """Set atom type (typically substructure) and (optionally) its f' and f''""" atomtype=atomtype.upper() if fp is not None and fpp is not None: try: fp=float(fp) fpp=float(fpp) except: common.Error("f' and f'' must be float ({0},{1} inputted for atom type {2})".format(fp,fpp,atomtype)) if not self.atomtype1: self.atomtype1=atomtype self.atomtypes[atomtype] = (fp,fpp) def SetAtomTypes(self,atomtypes,atomtype1=None): """Takes the input dictionary of atom types (and f',f'' if supplied) and sets it to this container""" if type(atomtypes) != dict: common.Error('SetAtomTypes() requires atom types dictionary as input') self.atomtypes=atomtypes if atomtype1: self.atomtype1=atomtype1 def GetAtomTypes(self, getlist=False): """Returns dictionary of atom types (typically substructure) with (f',f'') as dict. values by default or a list of atom types if 'getlist' is True """ if getlist: lst = self.atomtypes.keys() if self.atomtype1 and self.atomtypes[0]!=self.atomtype1: lst.remove(self.atomtype1) lst.insert(0, self.atomtype1) return lst else: return self.atomtypes def GetAtomType(self): """Returns the first (main) atom type """ return self.atomtype1 class exclude(data_container): """Excluded reflections container class""" description="'free' reflections to be excluded" supported_filetypes=('mtz',) deffiletype='mtz' types=('freeR','freebias') typ='freeR' col_list=['free',] label_default = { 'freeR': ['FREER'], \ 'freebias': ['BIASFREE'] } supported_attributes=('typ','file','xname','dname') class datafile(data_container): """Container class for various other, typically program specific, data storage""" # type is given by filetype: # ins = SHELX input parameters # xml = .xml file # mtz = .mtz file (if we are not interested in its crystallographic information, eg for mtz merging) # dluz = Refmac/Multicomb Luzzati D parameters file # par = ARP/wARP .par file description="other data - defined by file type" supported_filetypes=('ins','xml','mtz','dluz','par') deffiletype='ins' types=supported_filetypes typ=deffiletype supported_attributes=('typ','file') def __init__(self, filename=None, filetype=None, typ=None, *args, **kwargs): if filetype is None and typ: filetype=typ if typ is None and filetype: typ=filetype super(datafile,self).__init__(filename=filename,filetype=filetype,typ=typ,*args,**kwargs)