obdotnet_enumerator_ext.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using OpenBabel;

namespace OBDotNetExamples
{
    public static class OBDotNetExtensions
    {

        public static IEnumerable<OBAtom> Neighbors(this OBAtom atom)
        {
            return Bonds(atom).Select<OBBond, OBAtom>((bond) => bond.GetNbrAtom(atom));
        }

        public static bool Contains(this OBBond bond, OBAtom atom)
        {
            uint idx = atom.GetIdx();

            if (bond.GetBeginAtomIdx() == idx || bond.GetEndAtomIdx() == idx)
                return true;

            return false;
        }

        public static IEnumerable<OBBond> Bonds(this OBAtom atom)
        { 
            OBBond bond;
            int bondsFound = 0;

            for(int i = 0; i < atom.GetParent().NumBonds(); i++)
            {
                bond = atom.GetParent().GetBond(i);

                if (bond.Contains(atom))
                {
                    bondsFound++;
                    yield return bond;
                }

                else if (bondsFound == atom.GetValence())
                    yield break;
            }
            
        }

        public static IEnumerable<OBAtom> Atoms(this OBMol mol)
        {
           
            uint numAtoms = mol.NumAtoms();
            
            if(numAtoms == 0)
                yield return null;

            for(int i = 1; i <= numAtoms; i++)
                yield return mol.GetAtom(i);

        }

        public static IEnumerable<OBAtom> AtomsDFS(this OBMol mol)
        {
            return AtomsDFS(mol,1);
        }


        public static IEnumerable<OBAtom> AtomsDFS(this OBMol mol, int startingIndex)
        {
            if (startingIndex < 1)
                throw new ArgumentException("Atom indices must be greater than 0");

            if (startingIndex > mol.NumAtoms())
                throw new ArgumentOutOfRangeException("Index " + startingIndex + " is invalid. Structure contains " + mol.NumAtoms() + " atoms");

            bool[] visited = new bool[mol.NumAtoms()];
            visited[startingIndex - 1] = true;

            OBAtom current = mol.GetAtom(startingIndex);

            yield return current;

            Stack<OBAtom> path = new Stack<OBAtom>();

            foreach (OBAtom neighbor in current.Neighbors().Where(a => !visited[a.GetIdx() - 1]))
            {
                visited[neighbor.GetIdx() - 1] = true;
                path.Push(neighbor);

            }

            
            while (path.Count != 0)
            {
                
                current = path.Pop();
               
                yield return current;

                foreach (OBAtom neighbor in current.Neighbors().Where(a => !visited[a.GetIdx() - 1]))
                {
                    visited[neighbor.GetIdx() - 1] = true;
                    path.Push(neighbor);
               
                }
            }
            
        }

        public static IEnumerable<AtomDepthPair> AtomsBFSWDepth(this OBMol mol, int startingIndex)
        {

            if (startingIndex < 1)
                throw new ArgumentException("Atom indices must be greater than 0");

            if (startingIndex > mol.NumAtoms())
                throw new ArgumentOutOfRangeException("Index " + startingIndex + " is invalid. Structure contains " + mol.NumAtoms() + " atoms");

            
            bool[] visited = new bool[mol.NumAtoms()];
            visited[startingIndex - 1] = true;

            OBAtom first = mol.GetAtom(startingIndex);
            
            yield return new AtomDepthPair(first,1);

            Queue<AtomDepthPair> queue = new Queue<AtomDepthPair>();

            int depth = 2;

            foreach(OBAtom a in first.Neighbors().Where(a=>!visited[a.GetIdx()-1]))
            {
                queue.Enqueue(new AtomDepthPair(a, depth));
            }

            AtomDepthPair current;

            while (queue.Count != 0)
            {
                current = queue.Dequeue();
                visited[current.Atom.GetIdx() - 1] = true;
                      
                yield return current;

                foreach(OBAtom neighbor in current.Atom.Neighbors().Where(a=>!visited[a.GetIdx()-1]))
                    queue.Enqueue(new AtomDepthPair(neighbor,current.Depth+1));
            }


        }

        public static IEnumerable<OBAtom> AtomsBFS(this OBMol mol)
        {
            foreach(AtomDepthPair pair in AtomsBFSWDepth(mol, 1))
            {
                yield return pair.Atom;
            }
        }

        public static IEnumerable<OBBond> Bonds(this OBMol mol)
        {
            for (int i = 0; i < mol.NumBonds(); i++)
                yield return mol.GetBond(i);
        }

    }



    public struct AtomDepthPair
    {
        public int Depth;
        public OBAtom Atom;

        public AtomDepthPair(OBAtom atom, int depth)
        {
            Atom = atom;
            Depth = depth;
        }

    }
}