Skip to content
Snippets Groups Projects
BandStructureNormalizer.scala 10.62 KiB
/*
 * Copyright 2017-2018 Fawzi Mohamed
 * 
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

package eu.nomad_lab.normalizers

import eu.{ nomad_lab => lab }
import eu.nomad_lab.DefaultPythonInterpreter
import org.{ json4s => jn }
import scala.collection.breakOut
import eu.nomad_lab.normalize.ExternalNormalizerGenerator
import eu.nomad_lab.meta
import eu.nomad_lab.query
import eu.nomad_lab.resolve._
import eu.nomad_lab.ref.NomadUri
import eu.nomad_lab.h5.EmitJsonVisitor
import eu.nomad_lab.h5.SectionH5
import eu.nomad_lab.h5.H5EagerScanner
import eu.nomad_lab.parsers.ExternalParserWrapper
import eu.nomad_lab.JsonUtils
import scala.collection.mutable.StringBuilder

object BandStructureNormalizer extends ExternalNormalizerGenerator(
  name = "BandStructureNormalizer",
  info = jn.JObject(
    ("name" -> jn.JString("BandStructureNormalizer")) ::
      ("parserId" -> jn.JString("BandStructureNormalizer" + lab.BandStructureVersionInfo.version)) ::
      ("versionInfo" -> jn.JObject(
        ("nomadCoreVersion" -> jn.JObject(lab.NomadCoreVersionInfo.toMap.map {
          case (k, v) => k -> jn.JString(v.toString)
        }(breakOut): List[(String, jn.JString)])) ::
          (lab.BandStructureVersionInfo.toMap.map {
            case (key, value) =>
              (key -> jn.JString(value.toString))
          }(breakOut): List[(String, jn.JString)])
      )) :: Nil
  ),
  context = "calculation_context",
  filter = query.CompiledQuery(query.QueryExpression("section_k_band" /*" and not section_k_band_normalized" */ ), meta.KnownMetaInfoEnvs.all),
  cmd = Seq(DefaultPythonInterpreter.pythonExe(), "${envDir}/normalizers/band-structure/normalizer/normalizer-band-structure/normalizer_band_structure.py",
    "${contextUri}", "${archivePath}"),
  resList = Seq(
    "normalizer-band-structure/normalizer_band_structure.py",
    "normalizer-band-structure/setup_paths.py",
    "nomad_meta_info/public.nomadmetainfo.json",
    "nomad_meta_info/common.nomadmetainfo.json",
    "nomad_meta_info/meta_types.nomadmetainfo.json",
    "nomad_meta_info/stats.nomadmetainfo.json"
  ) ++ DefaultPythonInterpreter.commonFiles(),
  dirMap = Map(
    "normalizer-band-structure" -> "normalizers/band-structure/normalizer/normalizer-band-structure",
    "nomad_meta_info" -> "nomad-meta-info/meta_info/nomad_meta_info",
    "python" -> "python-common/common/python/nomadcore"
  ) ++ DefaultPythonInterpreter.commonDirMapping(),
  metaInfoEnv = lab.meta.KnownMetaInfoEnvs.stats

) {
  val trace: Boolean = false

  override def stdInHandler(context: ResolvedRef)(wrapper: ExternalParserWrapper)(pIn: java.io.OutputStream): Unit = {
    val out: java.io.Writer = if (trace)
      new java.io.BufferedWriter(new java.io.OutputStreamWriter(pIn));
    else
      null
    val stringBuilder = new StringBuilder
    def writeOut(s: String): Unit = {
      out.write(s)
      if (trace) stringBuilder ++= s
    }
    def flush(): Unit = {
      out.flush()
      if (trace) {
        logger.info(stringBuilder.result())
        stringBuilder.clear()
      }
    }
    writeOut("[")
    var isFirst = true
    try {
      context match {
        case Calculation(archiveSet, c) =>
          for (b <- c.sectionTable(Seq("section_run", "section_single_configuration_calculation", "section_k_band"))) {
            if (!isFirst)
              writeOut(",")
            else
              isFirst = false
            writeOut(s"""{
                          |  "context": ${JsonUtils.escapeString(b.toRef.toUriStr(archiveSet.objectKind))},
                          |  "section_k_band": """.stripMargin)
            val visitor = new EmitJsonVisitor(
              writeOut = writeOut
            )
            val scanner = new H5EagerScanner
            scanner.scanResolvedRef(Section(archiveSet, b), visitor)

            def writeOutRefE(refEName: String, refE: Option[Seq[Double]]): Unit = {
              refE match {
                case Some(e) =>
                  writeOut(s""",
                              |  "$refEName": ${e.mkString("[", ",", "]")}""".stripMargin)
                case None => ()
              }
            }

            b.parentSection.foreach { singleConf: SectionH5 =>
              val eFermi = singleConf.maybeValue("energy_reference_fermi").map(_.seqDoubleValue())
              writeOutRefE("energy_reference_fermi", eFermi)
              val eVbTop = singleConf.maybeValue("energy_reference_highest_occupied").map(_.seqDoubleValue())
              writeOutRefE("energy_reference_highest_occupied", eVbTop)
              var contextsToGiveBack: Seq[ResolvedRef] = Seq()
              try {
                // get DOS
                val dosInfo: Option[(Option[Seq[Double]], Option[Seq[Double]], SectionH5)] = b.parentSection match {
                  case None => None
                  case Some(singleConf) =>
                    // try locally
                    val dosColl = singleConf.subSectionCollection("section_dos")
                    if (dosColl.length > 0) {
                      if (dosColl.length > 1)
                        logger.warn(s"multiple dos in $singleConf, using first")
                      Some(eFermi, eVbTop, dosColl(0))
                    } else {
                      // look in archive
                      singleConf.maybeValue("single_configuration_calculation_to_system_ref").map(_.longValue) match {
                        case None => None
                        case Some(sysIdx) =>
                          singleConf.parentSection.map { run: SectionH5 =>
                            run.table.subSectionTable("section_system")(sysIdx)
                          } match {
                            case None => None
                            case Some(sys) =>
                              sys.maybeValue("configuration_raw_gid").map(_.stringValue) match {
                                case Some(confId) =>
                                  val possibleDosScanOp = new query.CollectUrisScanOp(
                                    context = "section_single_configuration_calculation",
                                    filter = query.CompiledQuery(query.QueryExpression("section_dos and configuration_raw_gid = \"$confId\""), meta.KnownMetaInfoEnvs.all),
                                    metaInfoEnv0 = None
                                  )
                                  val evaluator = new query.DirectArchiveScanner(
                                    archiveSet = archiveSet,
                                    archive = c.archive,
                                    scanOp = possibleDosScanOp
                                  )
                                  evaluator.run()
                                  val uris = possibleDosScanOp.uris
                                  evaluator.cleanup()
                                  if (uris.length > 0) {
                                    if (uris.length > 1) {
                                      logger.warn(s"found multiple DoS potentially matching to $b in archive: $uris, should check method more carefully, skipping for now...")
                                      None
                                    } else {
                                      val dosCtx = Resolver.resolveInArchiveSet(archiveSet, NomadUri(uris(0)).toRef)
                                      contextsToGiveBack = contextsToGiveBack :+ dosCtx
                                      dosCtx match {
                                        case Section(_, dosSingleConf) =>
                                          val dosEVbTop = dosSingleConf.maybeValue("energy_reference_highest_occupied").map(_.seqDoubleValue())
                                          val dosEFermi = dosSingleConf.maybeValue("energy_reference_fermi").map(_.seqDoubleValue())
                                          val dosCollection = dosSingleConf.subSectionCollection("section_dos")
                                          if (dosCollection.length > 0) {
                                            if (dosCollection.length > 1)
                                              logger.warn(s"multiple dos in $singleConf, using first")
                                            Some((dosEFermi, dosEVbTop, dosColl(0)))
                                          } else {
                                            logger.warn(s"query error, expected a section_dos in $dosSingleConf")
                                            None
                                          }
                                        case _ =>
                                          logger.warn(s"query error, expected valid section as result of query, not ${dosCtx.uriString}")
                                          None
                                      }
                                    }
                                  } else {
                                    None
                                  }
                                case None =>
                                  logger.warn(s"Could not find configuration_raw_gid for configuration of band structure $b, search of DoS failed.")
                                  None
                              }
                          }

                      }
                    }
                }
                dosInfo match {
                  case Some((dFermi, dVBTop, dos)) =>
                    writeOut(""",
                             |  "dos": {
                             |    "section_dos":""".stripMargin)

                    scanner.scanResolvedRef(Section(archiveSet, dos), visitor)
                    writeOutRefE("energy_reference_fermi", dFermi)
                    writeOutRefE("energy_reference_highest_occupied", dVBTop)
                    writeOut("\n  }")
                  case None => ()
                }
              } finally {
                for (ctx <- contextsToGiveBack)
                  ctx.giveBack()
              }
              writeOut("\n}")
            }
            flush()
          }
          writeOut("]\n")
          flush()
        case r =>
          throw new Exception(s"BandStructureNormalizer expected a calculation as context, but got $r")
      }
    } finally {
      out.close()
      pIn.close()
      wrapper.sendStatus = ExternalParserWrapper.SendStatus.Finished
    }
  }

}