Перенос молекулярной динамики на CUDA. Часть II: Суммирование по Эвальду

    В предыдущей статье мы обсудили основу метода молекулярной динамики, в том числе вычисление энергии и сил взаимодействия между частицами с заданными парными потенциалами. А что, если частицы обладают некоторым электрическим зарядом? Например, в том случае, если мы моделируем кристалл поваренной соли, состоящий из ионов Na+ и Cl-. Или водный раствор, содержащий те или иные ионы. В этом случае, кроме парных потенциалов типа Леннарда-Джонса между ионами действуют силы электростатического взаимодействия, т.е. закон Кулона. Энергия такого взаимодействия для пары частиц i-j равна:

    $E=C\frac{q_iq_j}{r_{ij}},$


    где q – заряд частицы, rij – расстояние между частицами, С – некоторая постоянная, зависящая от выбора единиц измерения. В системе СИ это — $\frac{1}{4\pi\epsilon_0}$, в СГС — 1, в моей программе (где энергия выражена в электронвольтах, расстояние в ангстремах, а заряд в элементарных зарядах) C примерно равно 14.3996.

    image

    Ну и что, скажете вы? Просто добавим соответствующее слагаемое в парный потенциал и готово. Однако, чаще всего в МД моделировании используют периодические граничные условия, т.е. моделируемая система со всех сторон окружена бесконечным количеством её виртуальных копий. В этом случае каждый виртуальный образ нашей системы будет взаимодействовать со всеми заряженными частицами внутри системы по закону Кулона. А поскольку Кулоновское взаимодействие убывает с расстоянием очень слабо (как 1/r), то отмахнуться от него так просто нельзя, сказав, что с такого-то расстояния мы его не вычисляем. Ряд вида 1/x расходится, т.е. его сумма, в принципе, может расти до бесконечности. И что же теперь, миску супа не солить? Убьёт электричеством?

    …можно не только солить суп, но и посчитать энергию Кулоновского взаимодействия в периодических граничных условиях. Такой метод был предложен Эвальдом ещё в 1921 году для расчета энергии ионного кристалла (можно также посмотреть в википедии). Суть метода заключатся в экранировании точечных зарядов и последующим вычетом функции экранирования. При этом часть электростатического взаимодействия сводится к короткойдействующему и его можно просто обрезать стандартным способом. Оставшаяся дальнодействующая часть эффективно суммируется в пространстве Фурье. Опуская вывод, который можно посмотреть в статье Блинова или в той же книге Френкеля и Смита сразу запишу решение, называемое суммой Эвальда:

    $E_{real} = С\sum\limits^{N-1}_i{\sum\limits^N_{j=i+1}{ \frac{q_iq_j}{r_{ij}}\operatorname{erfc}(\alpha r_{ij}) }}, $


    $E_{rec} = С\frac{2\pi}{V}\sum_{\vec{k}\ne0}(\exp(-k^2/4\alpha^2)/k^2 \sum_jq_i\exp(i\vec{k}\vec{r}_j)),$


    $E_{const} = -\frac{C}{V\alpha} \sum_iq_i^2, $


    где α – параметр, регулирующий соотношение вычислений в прямом и обратном пространствах, k – вектора в обратном пространстве по которым идёт суммирование, V – объём системы (в прямом пространстве). Первая часть (Ereal) является короткодействующей и вычисляется в том же цикле, что и другие парные потенциалы, смотри функцию real_ewald в предыдущей статье. Последний вклад (Eсonst) является поправкой на самовзаимодействие и часто называется «постоянной частью», поскольку не зависит от координат частиц. Её вычисление тривиально, поэтому мы остановимся только на второй части Эвальдовой суммы (Erec), суммировании в обратном пространстве. Естественно, во времена вывода Эвальда молекулярной динамики не было, кто впервые использовал этот метод в МД найти мне не удалось. Сейчас любая книга по МД содержит его изложение как некий золотой стандарт. К книге Аллена даже прилагается пример кода на фортране. К счастью, у меня остался код, написанный когда-то на С для последовательной версии, осталось только его распараллелить (я позволил себе опустить некоторые объявления переменных и другие несущественные детали):

    void ewald_rec()
    {
        int mmin = 0;
        int nmin = 1;
    
        // массивы где хранятся iexp(x[i] * kx[l]),
        double** elc;
        double** els;
        //... iexp(y[i] * ky[m]) и
        double** emc;
        double** ems;
        //... iexp(z[i] * kz[n]),
        double** enc;
        double** ens;
    
        // временные массивы для произведений iexp(x*kx)*iexp(y*ky)
        double* lmc;
        double* lms;
        // и для q[i] * iexp(x*kx)*iexp(y*ky)*iexp(z*kz)
        double* ckc;
        double* cks;
    
        // ПРЕДВАРИТЕЛЬНОЕ ЗАПОЛНЕНИЕ МАССИВОВ
        eng = 0.0;
        for (i = 0; i < Nat; i++)   // цикл по атомам
        {
            // emc/s[i][0] и enc/s[i][0] уже заполнены на этапе инициализации
            // в массив elc/s нужно обновить, см. далее
            elc[i][0] = 1.0;
            els[i][0] = 0.0;
    
            // iexp(kr)
            sincos(twopi * xs[i] * ra, els[i][1], elc[i][1]);
            sincos(twopi * ys[i] * rb, ems[i][1], emc[i][1]);
            sincos(twopi * zs[i] * rc, ens[i][1], enc[i][1]);
        }
    
        // заполняем следующие элементы массива emc/s[i][l] = iexp(y[i]*ky[l]) итеративно, используя комплексное умножение
        for (l = 2; l < ky; l++)
            for (i = 0; i < Nat; i++)
            {
                emc[i][l] = emc[i][l - 1] * emc[i][1] - ems[i][l - 1] * ems[i][1];
                ems[i][l] = ems[i][l - 1] * emc[i][1] + emc[i][l - 1] * ems[i][1];
            }
    
        // заполняем следующие элементы массива enc/s[i][l] = iexp(z[i]*kz[l]) итеративно, используя комплексное умножение
        for (l = 2; l < kz; l++)
            for (i = 0; i < Nat; i++)
            {
                enc[i][l] = enc[i][l - 1] * enc[i][1] - ens[i][l - 1] * ens[i][1];
                ens[i][l] = ens[i][l - 1] * enc[i][1] + enc[i][l - 1] * ens[i][1];
            }
    
        // ГЛАВНЫЙ ЦИКЛ ПО ВСЕМ K-ВЕКТОРАМ:
        for (l = 0; l < kx; l++)
        {
            rkx = l * twopi * ra; 
            // записываем exp(ikx[l]) в ikx[0] для сохранения памяти
            if (l == 1)
                for (i = 0; i < Nat; i++)
                {
                    elc[i][0] = elc[i][1];
                    els[i][0] = els[i][1];
                }
            else if (l > 1)
                for (i = 0; i < Nat; i++)
                {
                    // iexp(kx[0]) = iexp(kx[0]) * iexp(kx[1])
                    x = elc[i][0];
                    elc[i][0] = x * elc[i][1] - els[i][0] * els[i][1];
                    els[i][0] = els[i][0] * elc[i][1] + x * els[i][1];
                }
    
            for (m = mmin; m < ky; m++)
            {
                rky = m * twopi * rb;
                // заполняем временный массив произведением iexp(kx*x[i]) * iexp(ky*y[i])
                if (m >= 0)
                    for (i = 0; i < Nat; i++)
                    {
                        lmc[i] = elc[i][0] * emc[i][m] - els[i][0] * ems[i][m];
                        lms[i] = els[i][0] * emc[i][m] + ems[i][m] * elc[i][0];
                    }
                else // для отрицательных значений m используем комплексное сопряжение:
                    for (i = 0; i < Nat; i++)
                    {
                        lmc[i] = elc[i][0] * emc[i][-m] + els[i][0] * ems[i][-m];
                        lms[i] = els[i][0] * emc[i][-m] - ems[i][-m] * elc[i][0];
                    }
    
                for (n = nmin; n < kz; n++)
                {
                    rkz = n * twopi * rc;
                    rk2 = rkx * rkx + rky * rky + rkz * rkz;
                    if (rk2 < rkcut2) // используем радиус обрезания
                    {
                        // вычисляем сумму (q[i]*iexp(kr[k]*r[i])) - зарядовую плотность
                        sumC = 0; sumS = 0;
                        if (n >= 0)
                            for (i = 0; i < Nat; i++)
                            {
                                //считываем заряд частицы
                                ch = charges[types[i]].charge;
    
                                ckc[i] = ch * (lmc[i] * enc[i][n] - lms[i] * ens[i][n]);
                                cks[i] = ch * (lms[i] * enc[i][n] + lmc[i] * ens[i][n]);
    
                                sumC += ckc[i];
                                sumS += cks[i];
                            }
                        else // для отрицательных индексов используем комплексное сопряжение:
                            for (i = 0; i < Nat; i++)
                            {
                                //считываем заряд частицы
                                ch = charges[types[i]].charge;
    
                                ckc[i] = ch * (lmc[i] * enc[i][-n] + lms[i] * ens[i][-n]);
                                cks[i] = ch * (lms[i] * enc[i][-n] - lmc[i] * ens[i][-n]);
    
                                sumC += ckc[i];
                                sumS += cks[i];
                            }
    
                        //наконец вычисляем энергию и силы
                        akk = exp(rk2 * elec->mr4a2) / rk2;
                        eng += akk * (sumC * sumC + sumS * sumS);
    
                        for (i = 0; i < Nat; i++)
                        {
                            x = akk * (cks[i] * sumC - ckc[i] * sumS) * C * twopi * 2 * rvol;
                            fxs[i] += rkx * x;
                            fys[i] += rky * x;
                            fzs[i] += rkz * x;
                        }
    
                    }
                } // end n-loop (over kz-vectors)
    
                nmin = 1 - kz;
    
            } // end m-loop (over ky-vectors)
    
            mmin = 1 - ky;
    
        }  // end l-loop (over kx-vectors)
    
    
       engElec2 = eng * С * twopi * rvol;
    }
    

    Пара пояснений к коду: функция считает комплексную экспоненту (в комментариях к коду она обозначена iexp, чтобы убрать мнимую единицу из скобок) от векторного произведения k-вектора на радиус-вектор частицы для всех k-векторов и для всех частиц. Эта экспонента домножается на заряд частицы. Далее вычисляется сумма таких произведений по всем частицам (внутренняя сумма в формуле для Erec), у Френкеля она называется зарядовой плотностью, а у Блинова — структурным фактором. Ну а далее, на основании этих структурных факторов вычисляется энергия и силы, действующие на частицы. Компоненты k-векторов (2π*l/a, 2π*m/b, 2π*n/c) характеризуются тройкой целых чисел l, m и n, которые и пробегаются в циклах до заданных пользователем пределов. Параметры a, b и c – это размеры моделируемой системы в измерениях x, y и z соответственно (вывод верен для системы с геометрией прямоугольного параллелепипеда). В коде 1/a, 1/b и 1/с соответствуют переменным ra, rb и rc. Массивы под каждую величину представлены в двух экземплярах: под действительную и мнимую части. Каждый следующий k-вектор в одном измерении получается итеративно из предыдущего путем комплексного умножения предыдущего на единичный, чтобы каждый раз не считать синус с косинусом. Массивы emc/s и enc/s заполняются для всех m и n, соответственно, а массив elc/s значение для каждого l>1 помещает в нулевой индекс по l в целях экономии памяти.

    В целях распараллеливания выгодно «вывернуть» порядок циклов так, чтобы внешний цикл пробегался по частицам. И тут мы видим проблему – распараллелить эту функцию можно только до вычисления суммы по всем частицам (зарядовой плотности). Дальнейшие вычисления опираются на эту сумму, а она будет рассчитана только когда все потоки закончат работу, поэтому придётся разбить эту функцию на две. Первая вычисляет считает зарядовую плотность, а вторая – энергию и силы. Замечу, что во второй функции снова потребуется величина qiiexp(kr) для каждой частицы и для каждого k-вектора, вычисленная на предыдущем этапе. И тут есть два подхода: либо пересчитать её заново, либо запомнить.

    Первый вариант требует больше времени, второй – больше памяти (количество частиц * количество k-векторов * sizeof(float2)). Я остановился на втором варианте:

    __global__ void recip_ewald(int atPerBlock, int atPerThread, cudaMD* md)
    // calculate reciprocal part of Ewald summ
    // the first part : summ (qiexp(kr)) evaluation
    {
        int i;      // for atom loop
        int ik;     // index of k-vector
        int l, m, n;
        int mmin = 0;
        int nmin = 1;
        float tmp, ch;
        float rkx, rky, rkz, rk2;   // component of rk-vectors
    
        int nkx = md->nk.x;
        int nky = md->nk.y;
        int nkz = md->nk.z;
        
        // arrays for keeping iexp(k*r) Re and Im part
        float2 el[2];
        float2 em[NKVEC_MX];
        float2 en[NKVEC_MX];
    
        float2 sums[NTOTKVEC];          // summ (q iexp (k*r)) for each k
        extern __shared__ float2 sh_sums[];     // the same in shared memory
    
        float2 lm;     // temp var for keeping el*em
        float2 ck;     // temp var for keeping q * el * em * en (q iexp (kr))
    
        // invert length of box cell
        float ra = md->revLeng.x;
        float rb = md->revLeng.y;
        float rc = md->revLeng.z;
    
        if (threadIdx.x == 0)
            for (i = 0; i < md->nKvec; i++)
                sh_sums[i] = make_float2(0.0f, 0.0f);
        __syncthreads();
    
        for (i = 0; i < md->nKvec; i++)
            sums[i] = make_float2(0.0f, 0.0f);
    
        int id0 = blockIdx.x * atPerBlock + threadIdx.x * atPerThread;
        int N = min(id0 + atPerThread, md->nAt);
    
        ik = 0;
        for (i = id0; i < N; i++)
        {
            // save charge
            ch = md->specs[md->types[i]].charge;
    
            el[0] = make_float2(1.0f, 0.0f);    // .x - real part (or cos) .y - imagine part (or sin)
            
            em[0] = make_float2(1.0f, 0.0f);
            en[0] = make_float2(1.0f, 0.0f);
    
            // iexp (ikr)
            sincos(d_2pi * md->xyz[i].x * ra, &(el[1].y), &(el[1].x));
            sincos(d_2pi * md->xyz[i].y * rb, &(em[1].y), &(em[1].x));
            sincos(d_2pi * md->xyz[i].z * rc, &(en[1].y), &(en[1].x));
    
            // fil exp(iky) array by complex multiplication
            for (l = 2; l < nky; l++)
            {
                 em[l].x = em[l - 1].x * em[1].x - em[l - 1].y * em[1].y;
                 em[l].y = em[l - 1].y * em[1].x + em[l - 1].x * em[1].y;
            }
    
            // fil exp(ikz) array by complex multiplication
            for (l = 2; l < nkz; l++)
            {
                 en[l].x = en[l - 1].x * en[1].x - en[l - 1].y * en[1].y;
                 en[l].y = en[l - 1].y * en[1].x + en[l - 1].x * en[1].y;
            }
    
            // MAIN LOOP OVER K-VECTORS:
            for (l = 0; l < nkx; l++)
            {
                rkx = l * d_2pi * ra; 
                
                // move exp(ikx[l]) to ikx[0] for memory saving (ikx[i>1] are not used)
                if (l == 1)
                    el[0] = el[1];
                else if (l > 1)
                    {
                        // exp(ikx[0]) = exp(ikx[0]) * exp(ikx[1])
                        tmp = el[0].x;
                        el[0].x = tmp * el[1].x - el[0].y * el[1].y;
                        el[0].y = el[0].y * el[1].x + tmp * el[1].y;
                    }
    
                //ky - loop:
                for (m = mmin; m < nky; m++)
                {
                    rky = m * d_2pi * rb;
    
                    //set temporary variable lm = e^ikx * e^iky
                    if (m >= 0)
                    {
                            lm.x = el[0].x * em[m].x - el[0].y * em[m].y; 
                            lm.y = el[0].y * em[m].x + em[m].y * el[0].x;
                    }
                    else // for negative ky give complex adjustment to positive ky:
                    {
                            lm.x = el[0].x * em[-m].x + el[0].y * em[-m].y;
                            lm.y = el[0].y * em[-m].x - em[-m].x * el[0].x;
                    }
    
                    //kz - loop:
                    for (n = nmin; n < nkz; n++)
                    {
                        rkz = n * d_2pi * rc;
                        rk2 = rkx * rkx + rky * rky + rkz * rkz;
                        if (rk2 < md->rKcut2) // cutoff
                        {
                            // calculate summ[q iexp(kr)]   (local part)
                            if (n >= 0)
                             {
                                    ck.x = ch * (lm.x * en[n].x - lm.y * en[n].y);
                                    ck.y = ch * (lm.y * en[n].x + lm.x * en[n].y);
                             }
                            else // for negative kz give complex adjustment to positive kz:
                             {
                                    ck.x = ch * (lm.x * en[-n].x + lm.y * en[-n].y);
                                    ck.y = ch * (lm.y * en[-n].x - lm.x * en[-n].y);
                            }
                            sums[ik].x += ck.x;
                            sums[ik].y += ck.y;
                            
                            // save qiexp(kr) for each k for each atom:
                            md->qiexp[i][ik] = ck;
                            ik++;
                        }
                    } // end n-loop (over kz-vectors)
    
                    nmin = 1 - nkz;
    
                } // end m-loop (over ky-vectors)
                mmin = 1 - nky;
            }  // end l-loop (over kx-vectors)
        } // end loop by atoms
    
        // save sum into shared memory
        for (i = 0; i < md->nKvec; i++)
        {
            atomicAdd(&(sh_sums[i].x), sums[i].x);
            atomicAdd(&(sh_sums[i].y), sums[i].y);
        }
        __syncthreads();
    
        //...and to global
        int step = ceil((double)md->nKvec / (double)blockDim.x);
        id0 = threadIdx.x * step;
        N = min(id0 + step, md->nKvec);
        for (i = id0; i < N; i++)
        {
            atomicAdd(&(md->qDens[i].x), sh_sums[i].x);
            atomicAdd(&(md->qDens[i].y), sh_sums[i].y);
        }
    }
    // end 'ewald_rec' function
    
    __global__ void ewald_force(int atPerBlock, int atPerThread, cudaMD* md)
    // calculate reciprocal part of Ewald summ
    // the second part : enegy and forces
    {
        int i;      // for atom loop
        int ik;     // index of k-vector
        float tmp;
    
        // accumulator for force components
        float3 force;
    
        // constant factors for energy and force
        float eScale = md->ewEscale;
        float fScale = md->ewFscale;
    
        int id0 = blockIdx.x * atPerBlock + threadIdx.x * atPerThread;
        int N = min(id0 + atPerThread, md->nAt);
        for (i = id0; i < N; i++)
        {
            force = make_float3(0.0f, 0.0f, 0.0f);
    
            // summ by k-vectors
            for (ik = 0; ik < md->nKvec; ik++)
            {
                tmp = fScale * md->exprk2[ik] * (md->qiexp[i][ik].y * md->qDens[ik].x - md->qiexp[i][ik].x * md->qDens[ik].y);
    
                force.x += tmp * md->rk[ik].x;
                force.y += tmp * md->rk[ik].y;
                force.z += tmp * md->rk[ik].z;
            }
    
            md->frs[i].x += force.x;
            md->frs[i].y += force.y;
            md->frs[i].z += force.z;
        } // end loop by atoms
    
    
        // one block calculate energy
        if (blockIdx.x == 0)
            if (threadIdx.x == 0)
            {
                for (ik = 0; ik < md->nKvec; ik++)
                    md->engCoul2 += eScale * md->exprk2[ik] * (md->qDens[ik].x * md->qDens[ik].x + md->qDens[ik].y * md->qDens[ik].y);
            }
    
    }
    // end 'ewald_force' function
    

    Надеюсь, вы мне простите, что я оставил комментарии на английском, код практически повторяет последовательную версию. Код даже стал читабельнее за счет того, что массивы потеряли одно измерение: elc/s[i][l], emc/s[i][m] и enc/s[i][n] превратились в одномерные el, em и en, массивы lmc/s и ckc/s – в переменные lm и ck (пропала мерность по частицам, поскольку отпала необходимость хранить это для каждой частицы, промежуточный результат накапливается в shared memory). К сожалению, тут же возникла и проблема: массивы em и en пришлось задать статическими, чтобы не использовать глобальную память и не выделять память динамически каждый раз. Количество элементов в них определяется директивой NKVEC_MX (максимальное количество k-векторов по одному измерению) на этапе компиляции, а runtime используются только первые nky/z элементов. Кроме того, появился сквозной индекс по всем k-векторам и аналогичная директива, ограничивающая общее количество этих векторов NTOTKVEC. Проблема возникнет, если пользователю понадобится больше k-векторов, чем определено директивами. Для вычисления энергии предусмотрен блок с нулевым индексом, поскольку неважно какой именно блок выполнит этот расчет и в каком порядке. Замечу, что величина, вычисляемая в переменной akk серийного кода зависит только от размеров моделируемой системы и может быть посчитана на этапе инициализации, в моей реализации она хранится в массиве md->exprk2[] для каждого k-вектора. Аналогичным образом компоненты k-векторов берутся из массива md->rk[]. Тут может быть надо было использовать готовые функции FFT, раз уж метод основан на нём, но я так и не сообразил, как это сделать.

    Ну а теперь попробуем что-нибудь посчитать, да тот же содиум хлорайд. Возьмём 2 тысячи ионов натрия и столько же хлора. Заряды зададим целыми, а парные потенциалы возьмём, например, из этой работы. Стартовую конфигурацию зададим случайно и слегка перемешаем её, рисунок 2а. Объём системы выберем так, чтобы он соответствовал плотности поваренной соли при комнатной температуре (2,165 г/см3). Запустим все это на небольшое время (10’000 шагов по 5 фемтосекунд) с наивным учетом электростатики по закону Кулона и используя суммирование по Эвальду. Результирующие конфигурации приведены на рисунках 2б и 2в, соответственно. Вроде бы в случае с Эвальдом система чуть больше упорядочилась чем без него. Важно также, что флуктуации полной энергии с применением суммирования существенно уменьшились.


    Рисунок 2. Начальная конфигурация системы NaCl (a) и после 10’000 шагов интегрирования: наивным способом (б) и со схемой Эвальда (в).

    Вместо заключения


    Замечу, что структура, получаемая на рисунке, не соответствует кристаллической решетки NaCl, а скорее – решетке ZnS, но это уже претензия к парным потенциалам. Учет же электростатики очень важен для молекулярно-динамического моделирования. Считается, что именно электростатическое взаимодействие ответственно за образование кристаллических решёток, поскольку действует на больших расстояниях. Правда с этой позиции сложно объяснить как при охлаждении кристаллизуются такие вещества как аргон.

    Кроме упомянутого метода Эвальда, есть ещё и другие способы учета электростатики, смотрите, например, этот обзор.
    AdBlock похитил этот баннер, но баннеры не зубы — отрастут

    Подробнее
    Реклама

    Комментарии 4

      0
      Спасибо Вам, чудесная статья!
        0
        Благодарю
        0
        Спасибо. Очень круто!
        Буду признателен за мнение специалиста о перспективах создания программы для макетирования электронных оболочек атомов и молекул (электронной молекулярной динамики) на основе теории Кеннета Снельсона:
        habr.com/ru/post/374003
          0
          Статью прочитал, но не понял, что вы хотите.
          Там же только качественная модель. Ни какие свойства считать, ни как считать — ничего нет.

        Только полноправные пользователи могут оставлять комментарии. Войдите, пожалуйста.

        Самое читаемое