Python snl.StructureNL类代码示例

 def run_task(self, fw_spec):
     # Read structure from previous relaxation
     relaxed_struct = fw_spec['output']['crystal']
     # Generate deformed structures
     d_struct_set = DeformedStructureSet(relaxed_struct, ns=0.06)
     wf=[]
     for i, d_struct in enumerate(d_struct_set.def_structs):
         fws=[]
         connections={}
         f = Composition(d_struct.formula).alphabetical_formula
         snl = StructureNL(d_struct, 'Joseph Montoya <[email protected]>', 
                           projects=["Elasticity"])
         tasks = [AddSNLTask()]
         snl_priority = fw_spec.get('priority', 1)
         spec = {'task_type': 'Add Deformed Struct to SNL database', 
                 'snl': snl.as_dict(), 
                 '_queueadapter': QA_DB, 
                 '_priority': snl_priority}
         if 'snlgroup_id' in fw_spec and isinstance(snl, MPStructureNL):
             spec['force_mpsnl'] = snl.as_dict()
             spec['force_snlgroup_id'] = fw_spec['snlgroup_id']
             del spec['snl']
         fws.append(Firework(tasks, spec, 
                             name=get_slug(f + '--' + spec['task_type']), 
                             fw_id=-1000+i*10))
         connections[-1000+i*10] = [-999+i*10]
         spec = snl_to_wf._snl_to_spec(snl, 
                                       parameters={'exact_structure':True})
         spec = update_spec_force_convergence(spec)
         spec['deformation_matrix'] = d_struct_set.deformations[i].tolist()
         spec['original_task_id'] = fw_spec["task_id"]
         spec['_priority'] = fw_spec['_priority']*2
         #Turn off dupefinder for deformed structure
         del spec['_dupefinder']
         spec['task_type'] = "Optimize deformed structure"
         fws.append(Firework([VaspWriterTask(), SetupElastConstTask(),
                              get_custodian_task(spec)], 
                             spec, 
                             name=get_slug(f + '--' + spec['task_type']), 
                             fw_id=-999+i*10))
         
         priority = fw_spec['_priority']*3
         spec = {'task_type': 'VASP db insertion', 
                 '_priority': priority,
                 '_allow_fizzled_parents': True, 
                 '_queueadapter': QA_DB, 
                 'elastic_constant':"deformed_structure", 
                 'clean_task_doc':True,
                 'deformation_matrix':d_struct_set.deformations[i].tolist(), 
                 'original_task_id':fw_spec["task_id"]}
         fws.append(Firework([VaspToDBTask()], 
                             spec, 
                             name=get_slug(f + '--' + spec['task_type']), 
                             fw_id=-998+i*10))
         connections[-999+i*10] = [-998+i*10]
         wf.append(Workflow(fws, connections))
     return FWAction(additions=wf)

def structure_to_mock_job(structure):
    # Needs at least one author. This is for a mock job, so can put whatever.
    snl = StructureNL(structure, [{"name": "Saurabh Bajaj", "email": "[email protected]"},
                                  {"name": "Anubhav Jain", "email": "[email protected]"}])
    job = snl.as_dict()
    if 'is_valid' not in job:
        job.update(get_meta_from_structure(snl.structure))
    sorted_structure = snl.structure.get_sorted_structure()
    job.update(sorted_structure.as_dict())
    return job

    def test_to_from_dict(self):
        # no complicated objects in the 'data' or 'nodes' field
        a = StructureNL(self.s, self.hulk, ['test_project'], self.pmg,
                        ['remark1'], {"_my_data": "string"},
                        [self.valid_node, self.valid_node2])
        b = StructureNL.from_dict(a.as_dict())
        self.assertEqual(a, b)
        # complicated objects in the 'data' and 'nodes' field
        complicated_node = {
            "name": "complicated node",
            "url": "www.complicatednodegoeshere.com",
            "description": {
                "structure": self.s2
            }
        }
        a = StructureNL(self.s, self.hulk, ['test_project'], self.pmg,
                        ['remark1'], {"_my_data": {
                            "structure": self.s2
                        }}, [complicated_node, self.valid_node])
        b = StructureNL.from_dict(a.as_dict())
        self.assertEqual(
            a, b, 'to/from dict is broken when object embedding is '
            'used! Apparently MontyEncoding is broken...')

        #Test molecule
        molnl = StructureNL(self.mol, self.hulk, references=self.pmg)
        b = StructureNL.from_dict(molnl.as_dict())
        self.assertEqual(molnl, b)

    def run_task(self, fw_spec):
        sma = SNLMongoAdapter.auto_load()
        snl = StructureNL.from_dict(fw_spec['snl'])
        mpsnl, snlgroup_id, spec_group = sma.add_snl(snl)
        mod_spec = [{"_push": {"run_tags": "species_group={}".format(spec_group)}}] if spec_group else None

        return FWAction(update_spec={'mpsnl': mpsnl.to_dict, 'snlgroup_id': snlgroup_id}, mod_spec=mod_spec)

    def run_task(self, fw_spec):
        # pass-through option for when we start with an mpsnl and don't actually want to add
        if 'force_mpsnl' in fw_spec and 'force_snlgroup_id' in fw_spec:
            print 'USING FORCED MPSNL'
            return FWAction(update_spec={'mpsnl': fw_spec['force_mpsnl'], 'snlgroup_id': fw_spec['force_snlgroup_id']})

        sma = SNLMongoAdapter.auto_load()
        snl = StructureNL.from_dict(fw_spec['snl'])
        mpsnl, snlgroup_id = sma.add_snl(snl)

        return FWAction(update_spec={'mpsnl': mpsnl.to_dict, 'snlgroup_id': snlgroup_id})

    def run_task(self, fw_spec):
        # get the SNL mongo adapter
        sma = SNLMongoAdapter.auto_load()

        # get the SNL
        snl = StructureNL.from_dict(fw_spec['snl'])

        # add snl
        mpsnl, snlgroup_id = sma.add_snl(snl)

        return FWAction(update_spec={'mpsnl': mpsnl.to_dict, 'snlgroup_id': snlgroup_id})

    def from_snl(snl, snl_id, sg_num, sg_symbol, hall, xtal_system, lattice_type):
        # make a copy of SNL
        snl2 = StructureNL.from_dict(snl.to_dict)
        if '_materialsproject' not in snl2.data:
            snl2.data['_materialsproject'] = {}
        snl2.data['_materialsproject']['snl_id'] = snl_id
        snl2.data['_materialsproject']['sg_num'] = sg_num
        snl2.data['_materialsproject']['sg_symbol'] = sg_symbol
        snl2.data['_materialsproject']['hall'] = hall
        snl2.data['_materialsproject']['xtal_system'] = xtal_system
        snl2.data['_materialsproject']['lattice_type'] = lattice_type

        return MPStructureNL.from_dict(snl2.to_dict)

    def submit_structures(self, structures, authors, projects=None,
                             references='', remarks=None, data=None,
                             histories=None, created_at=None):
        """
        Submits a list of structures to the Materials Project as SNL files.
        The argument list mirrors the arguments for the StructureNL object,
        except that a list of structures with the same metadata is used as an
        input.

        .. note::

            As of now, this MP REST feature is open only to a select group of
            users. Opening up submissions to all users is being planned for
            the future.

        Args:
            structures:
                A list of Structure objects
            authors:
                *List* of {"name":'', "email":''} dicts,
                *list* of Strings as 'John Doe <[email protected]>',
                or a single String with commas separating authors
            projects:
                List of Strings ['Project A', 'Project B']. This applies to
                all structures.
            references:
                A String in BibTeX format. Again, this applies to all
                structures.
            remarks:
                List of Strings ['Remark A', 'Remark B']
            data:
                A list of free form dict. Namespaced at the root level with an
                underscore, e.g. {"_materialsproject":<custom data>}. The
                length of data should be the same as the list of structures
                if not None.
            histories:
                List of list of dicts - [[{'name':'', 'url':'',
                'description':{}}], ...] The length of histories should be the
                same as the list of structures if not None.
            created_at:
                A datetime object

        Returns:
            A list of inserted submission ids.
        """
        snl_list = StructureNL.from_structures(structures,
                    authors, projects, references, remarks, data,
                    histories, created_at)

        self.submit_snl(snl_list)

    def test_from_structures(self):
        s1 = Structure([[5, 0, 0], [0, 5, 0], [0, 0, 5]], ["Fe"], [[0, 0, 0]])
        s2 = Structure([[5, 0, 0], [0, 5, 0], [0, 0, 5]], ["Mn"], [[0, 0, 0]])
        remarks = ["unittest"]
        authors="Test User <[email protected]>"
        snl_list = StructureNL.from_structures([s1, s2], authors, remarks=remarks)

        self.assertEqual(len(snl_list), 2)
        snl1 = snl_list[0]
        snl2 = snl_list[1]
        self.assertEqual(snl1.remarks, remarks)
        self.assertEqual(snl2.remarks, remarks)
        self.assertEqual(snl1.authors, [Author.parse_author(authors)])
        self.assertEqual(snl2.authors, [Author.parse_author(authors)])

    def submit_new_workflow(self):
        # finds a submitted job, creates a workflow, and submits it to FireWorks
        job = self.jobs.find_and_modify({'state': 'SUBMITTED'}, {'$set': {'state': 'WAITING'}})
        if job:
            submission_id = job['submission_id']
            try:
                if 'snl_id' in job:
                    snl = MPStructureNL.from_dict(job)
                else:
                    snl = StructureNL.from_dict(job)
                if len(snl.structure.sites) > SubmissionProcessor.MAX_SITES:
                    self.sma.update_state(submission_id, 'REJECTED', 'too many sites', {})
                    print 'REJECTED WORKFLOW FOR {} - too many sites ({})'.format(
                        snl.structure.formula, len(snl.structure.sites))
                elif not job['is_valid']:
                    self.sma.update_state(submission_id, 'REJECTED',
                                          'invalid structure (atoms too close)', {})
                    print 'REJECTED WORKFLOW FOR {} - invalid structure'.format(
                        snl.structure.formula)
                elif len(set(NO_POTCARS) & set(job['elements'])) > 0:
                    self.sma.update_state(submission_id, 'REJECTED',
                                          'invalid structure (no POTCAR)', {})
                    print 'REJECTED WORKFLOW FOR {} - invalid element (No POTCAR)'.format(
                        snl.structure.formula)
                elif not job['is_ordered']:
                    self.sma.update_state(submission_id, 'REJECTED',
                                          'invalid structure (disordered)', {})
                    print 'REJECTED WORKFLOW FOR {} - invalid structure'.format(
                        snl.structure.formula)
                else:
                    snl.data['_materialsproject'] = snl.data.get('_materialsproject', {})
                    snl.data['_materialsproject']['submission_id'] = submission_id

                    # create a workflow
                    if "Elasticity" in snl.projects:
                        from mpworks.workflows.snl_to_wf_phonon import snl_to_wf_phonon
                        wf=snl_to_wf_phonon(snl, job['parameters'])
                    else:
                        wf = snl_to_wf(snl, job['parameters'])
                    self.launchpad.add_wf(wf)
                    print 'ADDED WORKFLOW FOR {}'.format(snl.structure.formula)
            except:
                self.jobs.find_and_modify({'submission_id': submission_id},
                                          {'$set': {'state': 'ERROR'}})
                traceback.print_exc()

            return submission_id

    def from_snl(snl, snl_id, sg_num, sg_symbol, hall, xtal_system, lattice_type, pointgroup):
        # make a copy of SNL
        snl2 = StructureNL.from_dict(snl.to_dict)
        if '_materialsproject' not in snl2.data:
            snl2.data['_materialsproject'] = {}

        snl2.data['_materialsproject']['snl_id'] = snl_id
        snl2.data['_materialsproject']['spacegroup'] = {}
        sg = snl2.data['_materialsproject']['spacegroup']
        sg['symbol'] = sg_symbol
        sg['number'] = sg_num
        sg['point_group'] = pointgroup
        sg['crystal_system'] = xtal_system
        sg['hall'] = hall
        sg['lattice_type'] = lattice_type

        return MPStructureNL.from_dict(snl2.to_dict)

    def from_snl(snl, snl_id, sg_num, sg_symbol, hall, xtal_system, lattice_type, pointgroup):
        # make a copy of SNL
        snl2 = StructureNL.from_dict(snl.as_dict())
        if "_materialsproject" not in snl2.data:
            snl2.data["_materialsproject"] = {}

        snl2.data["_materialsproject"]["snl_id"] = snl_id
        snl2.data["_materialsproject"]["spacegroup"] = {}
        sg = snl2.data["_materialsproject"]["spacegroup"]
        sg["symbol"] = sg_symbol
        sg["number"] = sg_num
        sg["point_group"] = pointgroup
        sg["crystal_system"] = xtal_system
        sg["hall"] = hall
        sg["lattice_type"] = lattice_type

        return MPStructureNL.from_dict(snl2.as_dict())

    def submit_new_workflow(self):
        # finds a submitted job, creates a workflow, and submits it to FireWorks
        job = self.jobs.find_and_modify({'state': 'submitted'}, {'$set': {'state': 'waiting'}})
        if job:
            submission_id = job['submission_id']
            try:
                snl = StructureNL.from_dict(job)
                snl.data['_materialsproject'] = snl.data.get('_materialsproject', {})
                snl.data['_materialsproject']['submission_id'] = submission_id

                # create a workflow
                wf = snl_to_wf(snl)
                self.launchpad.add_wf(wf)
                print 'ADDED WORKFLOW FOR {}'.format(snl.structure.formula)
            except:
                self.jobs.find_and_modify({'submission_id': submission_id}, {'$set': {'state': 'error'}})
                traceback.print_exc()

            return submission_id

def job_is_submittable(job):
    snl = StructureNL.from_dict(job)
    # mpworks.processors.process_submissions.SubmissionProcessor#submit_new_workflow
    max_sites = 200  # SubmissionProcessor.MAX_SITES above
    # from mpworks.workflows.wf_utils import NO_POTCARS
    no_potcars = ['Po', 'At', 'Rn', 'Fr', 'Ra', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr']
    if len(snl.structure.sites) > max_sites:
        print 'REJECTED WORKFLOW FOR {} - too many sites ({})'.format(
            snl.structure.formula, len(snl.structure.sites))
    elif not job['is_valid']:
        print 'REJECTED WORKFLOW FOR {} - invalid structure (atoms too close)'.format(
            snl.structure.formula)
    elif len(set(no_potcars) & set(job['elements'])) > 0:
        print 'REJECTED WORKFLOW FOR {} - invalid element (No POTCAR)'.format(
            snl.structure.formula)
    elif not job['is_ordered']:
        print 'REJECTED WORKFLOW FOR {} - invalid structure (disordered)'.format(
            snl.structure.formula)
    else:
        return True
    return False

    def run_task(self, fw_spec):
        print "sleeping 10s for Mongo"
        time.sleep(10)
        print "done sleeping"
        print "the gap is {}, the cutoff is {}".format(fw_spec["analysis"]["bandgap"], self.gap_cutoff)
        if fw_spec["analysis"]["bandgap"] >= self.gap_cutoff:
            static_dens = 90
            uniform_dens = 1000
            line_dens = 20
        else:
            static_dens = 450
            uniform_dens = 1500
            line_dens = 30

        if fw_spec["analysis"]["bandgap"] <= self.metal_cutoff:
            user_incar_settings = {"ISMEAR": 1, "SIGMA": 0.2}
        else:
            user_incar_settings = {}

        print "Adding more runs..."

        type_name = "GGA+U" if "GGA+U" in fw_spec["prev_task_type"] else "GGA"

        snl = StructureNL.from_dict(fw_spec["mpsnl"])
        f = Composition(snl.structure.composition.reduced_formula).alphabetical_formula

        fws = []
        connections = {}

        priority = fw_spec["_priority"]
        trackers = [
            Tracker("FW_job.out"),
            Tracker("FW_job.error"),
            Tracker("vasp.out"),
            Tracker("OUTCAR"),
            Tracker("OSZICAR"),
        ]
        trackers_db = [Tracker("FW_job.out"), Tracker("FW_job.error")]

        # run GGA static
        spec = fw_spec  # pass all the items from the current spec to the new
        spec.update(
            {
                "task_type": "{} static v2".format(type_name),
                "_queueadapter": QA_VASP_SMALL,
                "_dupefinder": DupeFinderVasp().to_dict(),
                "_priority": priority,
                "_trackers": trackers,
            }
        )
        fws.append(
            Firework(
                [
                    VaspCopyTask({"use_CONTCAR": True, "skip_CHGCAR": True}),
                    SetupStaticRunTask({"kpoints_density": static_dens, "user_incar_settings": user_incar_settings}),
                    get_custodian_task(spec),
                ],
                spec,
                name=get_slug(f + "--" + spec["task_type"]),
                fw_id=-10,
            )
        )

        # insert into DB - GGA static
        spec = {
            "task_type": "VASP db insertion",
            "_queueadapter": QA_DB,
            "_allow_fizzled_parents": True,
            "_priority": priority * 2,
            "_dupefinder": DupeFinderDB().to_dict(),
            "_trackers": trackers_db,
        }
        fws.append(Firework([VaspToDBTask()], spec, name=get_slug(f + "--" + spec["task_type"]), fw_id=-9))
        connections[-10] = -9

        # run GGA Uniform
        spec = {
            "task_type": "{} Uniform v2".format(type_name),
            "_queueadapter": QA_VASP,
            "_dupefinder": DupeFinderVasp().to_dict(),
            "_priority": priority,
            "_trackers": trackers,
        }
        fws.append(
            Firework(
                [
                    VaspCopyTask({"use_CONTCAR": False}),
                    SetupNonSCFTask({"mode": "uniform", "kpoints_density": uniform_dens}),
                    get_custodian_task(spec),
                ],
                spec,
                name=get_slug(f + "--" + spec["task_type"]),
                fw_id=-8,
            )
        )
        connections[-9] = -8

        # insert into DB - GGA Uniform
        spec = {
            "task_type": "VASP db insertion",
#.........这里部分代码省略.........