Add atom object name on the cmd.get_model() returned model (#380)

schrodinger · Aug 7, 2024 · d6b4762 · d6b4762
1 parent 03d7a7f
commit d6b4762
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 16 deletions.
diff --git a/layer2/CoordSet.cpp b/layer2/CoordSet.cpp
@@ -1,16 +1,16 @@
 
-/* 
+/*
 A* -------------------------------------------------------------------
 B* This file contains source code for the PyMOL computer program
-C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. 
+C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific.
 D* -------------------------------------------------------------------
 E* It is unlawful to modify or remove this copyright notice.
 F* -------------------------------------------------------------------
-G* Please see the accompanying LICENSE file for further information. 
+G* Please see the accompanying LICENSE file for further information.
 H* -------------------------------------------------------------------
 I* Additional authors of this source file include:
--* 
--* 
+-*
+-*
 -*
 Z* -------------------------------------------------------------------
 */
@@ -1054,8 +1054,8 @@ void CoordSetAtomToPDBStrVLA(PyMOLGlobals * G, char **charVLA, int *c,
 
 
 /*========================================================================*/
-PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const float *v,
-                                   const float *ref, int index, const double *matrix)
+PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, ObjectMolecule *obj,
+                                   const float *v, const float *ref, int index, const double *matrix)
 {
 #ifdef _PYMOL_NOPY
   return nullptr;
@@ -1072,6 +1072,7 @@ PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const fl
       RotateU(matrix, tmp_array);
     }
 
+    PConvStringToPyObjAttr(atom, "model", obj->Name);
     PConvFloat3ToPyObjAttr(atom, "coord", v);
     if(ref)
       PConvFloat3ToPyObjAttr(atom, "ref_coord", ref);
@@ -1130,6 +1131,7 @@ PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const fl
     PConvIntToPyObjAttr(atom, "id", ai->id);    /* not necc. unique */
     PConvIntToPyObjAttr(atom, "index", index + 1);      /* fragile */
 
+
 #ifdef _PYMOL_IP_PROPERTIES
 #endif
 

diff --git a/layer2/CoordSet.h b/layer2/CoordSet.h
@@ -1,16 +1,16 @@
 
-/* 
+/*
 A* -------------------------------------------------------------------
 B* This file contains source code for the PyMOL computer program
-C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. 
+C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific.
 D* -------------------------------------------------------------------
 E* It is unlawful to modify or remove this copyright notice.
 F* -------------------------------------------------------------------
-G* Please see the accompanying LICENSE file for further information. 
+G* Please see the accompanying LICENSE file for further information.
 H* -------------------------------------------------------------------
 I* Additional authors of this source file include:
--* 
--* 
+-*
+-*
 -*
 Z* -------------------------------------------------------------------
 */
@@ -122,7 +122,7 @@ struct CoordSet : CObjectState {
 
   pymol::vla<RefPosType> RefPos;
 
-  /* idea:  
+  /* idea:
      int start_atix, stop_atix <-- for discrete objects, we need
      something like this that would enable pymol to skip atoms not in the
      discrete state...question is: are these atoms sorted together right
@@ -202,8 +202,8 @@ bool CoordSetInsureOrthogonal(PyMOLGlobals * G,
     bool quiet=true);
 
 void CoordSetGetAverage(const CoordSet * I, float *v0);
-PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, const float *v,
-                                   const float *ref, int index, const double *matrix);
+PyObject *CoordSetAtomToChemPyAtom(PyMOLGlobals * G, AtomInfoType * ai, ObjectMolecule * obj,
+                                   const float *v, const float *ref, int index, const double *matrix);
 int CoordSetGetAtomVertex(const CoordSet * I, int at, float *v);
 int CoordSetGetAtomTxfVertex(const CoordSet * I, int at, float *v);
 int CoordSetSetAtomVertex(CoordSet * I, int at, const float *v);

diff --git a/layer3/MoleculeExporter.cpp b/layer3/MoleculeExporter.cpp
@@ -1820,7 +1820,7 @@ struct MoleculeExporterChemPy : public MoleculeExporter {
 
   void writeAtom() override {
     PyObject *atom = CoordSetAtomToChemPyAtom(G,
-        m_iter.getAtomInfo(), m_coord, getRefPtr(),
+        m_iter.getAtomInfo(), m_iter.obj, m_coord, getRefPtr(),
         m_iter.getAtm(), m_mat_full.ptr);
 
     if (atom) {

diff --git a/testing/tests/api/exporting.py b/testing/tests/api/exporting.py
@@ -348,6 +348,18 @@ def testPSEBulkExportImport(self, pse_export_version, pse_binary_dump):
             m2 = cmd.get_model()
             self.assertModelsAreSame(m1, m2)
 
+    def testGetModelObjectName(self):
+        cmd.load(self.datafile("1oky-frag.pdb"))
+        cmd.load(self.datafile('1rna.cif'))
+
+        m1 = cmd.get_model()
+        cnt = cmd.count_atoms('%1oky-frag')
+
+        self.assertEqual(m1.atom[0].model, '1oky-frag')
+        self.assertEqual(m1.atom[-1].model, '1rna')
+        self.assertEqual(m1.atom[cnt].model, '1rna')
+
+
     @testing.requires_version('2.1')
     def testMMTF(self):
         '''Styled MMTF export/import'''