Commit 1244af3a authored by Walid Ahmad's avatar Walid Ahmad
Browse files

Merge branch 'walid/pass_featurizer_kwargs' into walid/roberta-tokenizer

parents 81c50194 cd7215a1
Loading
Loading
Loading
Loading
+14 −10
Original line number Diff line number Diff line
@@ -161,7 +161,8 @@ class ComplexFeaturizer(Featurizer):

  def featurize(self,
                complexes: Iterable[Tuple[str, str]],
                log_every_n: int = 100) -> np.ndarray:
                log_every_n: int = 100,
                **kwargs) -> np.ndarray:
    """
    Calculate features for mol/protein complexes.

@@ -185,7 +186,7 @@ class ComplexFeaturizer(Featurizer):
      if idx % log_every_n == 0:
        logger.info("Featurizing datapoint %i" % idx)
      try:
        features.append(self._featurize(point))
        features.append(self._featurize(point, **kwargs))
        successes.append(idx)
      except:
        logger.warning(
@@ -208,7 +209,7 @@ class ComplexFeaturizer(Featurizer):

    return np.asarray(features)

  def _featurize(self, complex: Tuple[str, str]):
  def _featurize(self, complex: Tuple[str, str], **kwargs):
    """
    Calculate features for single mol/protein complex.

@@ -238,7 +239,7 @@ class MolecularFeaturizer(Featurizer):
  The subclasses of this class require RDKit to be installed.
  """

  def featurize(self, molecules, log_every_n=1000) -> np.ndarray:
  def featurize(self, molecules, log_every_n=1000, **kwargs) -> np.ndarray:
    """Calculate features for molecules.

    Parameters
@@ -282,7 +283,7 @@ class MolecularFeaturizer(Featurizer):
          new_order = rdmolfiles.CanonicalRankAtoms(mol)
          mol = rdmolops.RenumberAtoms(mol, new_order)

        features.append(self._featurize(mol))
        features.append(self._featurize(mol, **kwargs))
      except Exception as e:
        if isinstance(mol, Chem.rdchem.Mol):
          mol = Chem.MolToSmiles(mol)
@@ -320,7 +321,8 @@ class MaterialStructureFeaturizer(Featurizer):

  def featurize(self,
                structures: Iterable[Union[Dict[str, Any], PymatgenStructure]],
                log_every_n: int = 1000) -> np.ndarray:
                log_every_n: int = 1000,
                **kwargs) -> np.ndarray:
    """Calculate features for crystal structures.

    Parameters
@@ -351,7 +353,7 @@ class MaterialStructureFeaturizer(Featurizer):
      try:
        if isinstance(structure, Dict):
          structure = Structure.from_dict(structure)
        features.append(self._featurize(structure))
        features.append(self._featurize(structure, **kwargs))
      except:
        logger.warning(
            "Failed to featurize datapoint %i. Appending empty array" % idx)
@@ -383,8 +385,10 @@ class MaterialCompositionFeaturizer(Featurizer):
  installed.
  """

  def featurize(self, compositions: Iterable[str],
                log_every_n: int = 1000) -> np.ndarray:
  def featurize(self,
                compositions: Iterable[str],
                log_every_n: int = 1000,
                **kwargs) -> np.ndarray:
    """Calculate features for crystal compositions.

    Parameters
@@ -412,7 +416,7 @@ class MaterialCompositionFeaturizer(Featurizer):
        logger.info("Featurizing datapoint %i" % idx)
      try:
        c = Composition(composition)
        features.append(self._featurize(c))
        features.append(self._featurize(c, **kwargs))
      except:
        logger.warning(
            "Failed to featurize datapoint %i. Appending empty array" % idx)