Branch data Line data Source code
1 : : // *****************************************************************************
2 : : /*!
3 : : \file src/DiffEq/WrightFisher/WrightFisher.hpp
4 : : \copyright 2012-2015 J. Bakosi,
5 : : 2016-2018 Los Alamos National Security, LLC.,
6 : : 2019-2021 Triad National Security, LLC.
7 : : All rights reserved. See the LICENSE file for details.
8 : : \brief Wright-Fisher SDE
9 : : \details This file implements the time integration of a system of stochastic
10 : : differential equations (SDEs), whose invariant is the Dirichlet
11 : : distribution. For more details on the Wright-Fisher SDE, see
12 : : http://www.sciencedirect.com/science/article/pii/S0040580912001013.
13 : : */
14 : : // *****************************************************************************
15 : : #ifndef WrightFisher_h
16 : : #define WrightFisher_h
17 : :
18 : : #include <numeric>
19 : : #include <vector>
20 : : #include <iomanip>
21 : :
22 : : #include "WalkerBuildConfig.hpp"
23 : :
24 : : #ifdef HAS_MKL
25 : : #include <mkl_lapacke.h>
26 : : #else
27 : : #include <lapacke.h>
28 : : #endif
29 : :
30 : : #include "Macro.hpp"
31 : : #include "InitPolicy.hpp"
32 : : #include "WrightFisherCoeffPolicy.hpp"
33 : : #include "RNG.hpp"
34 : : #include "Particles.hpp"
35 : :
36 : : namespace walker {
37 : :
38 : : extern ctr::InputDeck g_inputdeck;
39 : : extern std::map< tk::ctr::RawRNGType, tk::RNG > g_rng;
40 : :
41 : : //! \brief Wright-Fisher SDE used polymorphically with DiffEq
42 : : //! \details The template arguments specify policies and are used to configure
43 : : //! the behavior of the class. The policies are:
44 : : //! - Init - initialization policy, see DiffEq/InitPolicy.h
45 : : //! - Coefficients - coefficients policy, see DiffEq/WrightFisherCoeffPolicy.h
46 : : template< class Init, class Coefficients >
47 [ - - ][ - - ]: 0 : class WrightFisher {
[ - - ][ - - ]
[ - - ][ - - ]
[ - - ][ - - ]
[ - - ][ - - ]
[ - - ][ - - ]
[ - - ][ - - ]
[ - - ][ - - ]
48 : :
49 : 0 : void print_matrix( const char *name, const double *mat, lapack_int n,
50 : : lapack_int m ) const
51 : : {
52 : : int i, j;
53 : 0 : printf("%s:\n", name);
54 [ - - ]: 0 : for(i=0; i<n; ++i) {
55 [ - - ]: 0 : for(j=0; j<m; ++j)
56 : 0 : printf("%10g ", mat[i*n+j]);
57 : 0 : printf("\n");
58 : : }
59 : 0 : printf("\n");
60 : 0 : }
61 : :
62 : : private:
63 : : using ncomp_t = kw::ncomp::info::expect::type;
64 : :
65 : : public:
66 : : //! \brief Constructor
67 : : //! \param[in] c Index specifying which system of Wright-Fisher SDEs to
68 : : //! construct. There can be multiple wright-fisher ... end blocks in a
69 : : //! control file. This index specifies which Wright-Fisher SDE system to
70 : : //! instantiate. The index corresponds to the order in which the
71 : : //! wright-fisher ... end blocks are given the control file.
72 : 0 : explicit WrightFisher( ncomp_t c ) :
73 : : m_c( c ),
74 : : m_depvar(
75 : : g_inputdeck.get< tag::param, tag::wrightfisher, tag::depvar >().at(c) ),
76 : : m_ncomp(
77 : : g_inputdeck.get< tag::component >().get< tag::wrightfisher >().at(c) ),
78 : : m_offset(
79 : 0 : g_inputdeck.get< tag::component >().offset< tag::wrightfisher >(c) ),
80 : 0 : m_rng( g_rng.at( tk::ctr::raw(
81 : : g_inputdeck.get< tag::param, tag::wrightfisher, tag::rng >().at(c) ) ) ),
82 : : m_omega(),
83 : : coeff(
84 : 0 : m_ncomp,
85 : : g_inputdeck.get< tag::param, tag::wrightfisher, tag::omega >().at(c),
86 [ - - ][ - - ]: 0 : m_omega )
[ - - ][ - - ]
[ - - ]
87 : : {
88 [ - - ][ - - ]: 0 : Throw( "Wright-Fisher diffusion matrix not yet implemented! See comments "
[ - - ][ - - ]
[ - - ][ - - ]
89 : : "in code for details." );
90 : : }
91 : : //! Initalize SDE, prepare for time integration
92 : : //! \param[in] stream Thread (or more precisely stream) ID
93 : : //! \param[in,out] particles Array of particle properties
94 : 0 : void initialize( [[maybe_unused]] int stream, tk::Particles& particles ) {
95 : : //! Set initial conditions using initialization policy
96 : : //Init::template
97 : : // init< tag::wrightfisher >
98 : : // ( g_inputdeck, m_rng, stream, particles, m_c, m_ncomp, m_offset );
99 : :
100 : : const auto npar = particles.nunk();
101 [ - - ]: 0 : for (auto p=decltype(npar){0}; p<npar; ++p) {
102 : : // Initialize the first m_ncomp (N-1) scalars
103 : : ncomp_t i;
104 [ - - ]: 0 : for (i=0; i<m_ncomp-1; ++i) {
105 : 0 : particles( p, i, m_offset ) =
106 : 0 : (1.0 + static_cast<tk::real>(i)) / static_cast<tk::real>(m_ncomp);
107 : : }
108 : : // Initialize the (N-1)th scalar from unit-sum
109 : 0 : tk::real& par = particles( p, i, m_offset );
110 : 0 : par = 1.0 - particles( p, 0, m_offset );
111 [ - - ]: 0 : for (i=1; i<m_ncomp-1; ++i) {
112 : 0 : par -= particles( p, i, m_offset );
113 : : }
114 : : }
115 : 0 : }
116 : :
117 : : //! \brief Advance particles according to the Wright-Fisher SDE
118 : : //! \param[in,out] particles Array of particle properties
119 : : //! \param[in] stream Thread (or more precisely stream) ID
120 : : //! \param[in] dt Time step size
121 : 0 : void advance( tk::Particles& particles,
122 : : int stream,
123 : : tk::real dt,
124 : : tk::real,
125 : : const std::map< tk::ctr::Product, tk::real >& )
126 : : {
127 : : // Compute sum of coefficients
128 : : const auto omega = std::accumulate( begin(m_omega), end(m_omega), 0.0 );
129 : : const auto npar = particles.nunk();
130 : :
131 : : #if defined(__clang__)
132 : : #pragma clang diagnostic push
133 : : #pragma clang diagnostic ignored "-Wvla"
134 : : #pragma clang diagnostic ignored "-Wvla-extension"
135 : : #elif defined(STRICT_GNUC)
136 : : #pragma GCC diagnostic push
137 : : #pragma GCC diagnostic ignored "-Wvla"
138 : : #endif
139 : :
140 [ - - ]: 0 : for (auto p=decltype(npar){0}; p<npar; ++p) {
141 : : // Need to build the square-root of the Wright-Fisher diffusion matrix:
142 : : // B_ij = y_i * ( delta_ij - y_j ). If the matrix is positive definite,
143 : : // the Cholesky decomposition would work, however, B_ij is only positive
144 : : // semi-definite, so Cholesky may fail and considering floating-point
145 : : // errors it will definitely fail at some point.
146 : : //
147 : : // A stable square-root for B_ij is not yet implemented. Here is some
148 : : // details and points for further development:
149 : : //
150 : : // An interesting fact from http://math.stackexchange.com/a/332465, on
151 : : // gauging the eigen values of a matrix:
152 : : //
153 : : // "Choosing each diagonal entry to be greater than the sum of the
154 : : // absolute values of the other entries in the same row will immediately
155 : : // imply that all of the eigenvalues of A are positive, and therefore
156 : : // that A is positive definite."
157 : : //
158 : : // The WF diffusion matrix seems to be positive semi-definite with at
159 : : // least one very small eigenvalue O(10e-18). The row and column sums
160 : : // are zero. It does not appear that the matrix has negative
161 : : // eigenvalues.
162 : : //
163 : : // Using Cholesky decomposition to compute the square-root may fail.
164 : : // Other methods, such as diagonalizing the matrix or the LDL
165 : : // decomposition, may work, but have not tried. The former requires
166 : : // finding the eigenvalues, with e.g., LAPACK's DSPEVD, thereby
167 : : // diagonalizing the matrix, and taking the square-root of the diagonal
168 : : // elements to find the square-root of the matrix. See
169 : : // http://en.wikipedia.org/wiki/Square_root_of_a_matrix. This may be
170 : : // okay, but probably overkill, as we don't really need the eigenvalues,
171 : : // only the square-root of the matrix.
172 : : //
173 : : // The latter (LDL) is similar to Cholesky, but does not involve square
174 : : // root. See http://en.wikipedia.org/wiki/Cholesky_decomposition. See
175 : : // also Golub van Loan: Sec.4.1.2 LDL. It appears that LDL would compare
176 : : // in perfromance to Cholesky and would work even on large indefinite
177 : : // matrices. LAPACK probably has LDL.
178 : : //
179 : : // I'm putting this aside for now. The next step should probably be
180 : : // idintifying the LAPACK for LDL and use it to advance the WF system
181 : : // using LD^{1/2}.
182 : : //
183 : : // See also:
184 : : // http://en.wikipedia.org/wiki/Positive-definite_matrix
185 : : // file:///opt/intel/composerxe/Documentation/en_US/mkl/mklman/index.htm
186 : : // http://icl.cs.utk.edu/lapack-forum/archives/lapack/msg01497.html
187 : : // http://yarchive.net/comp/sqrtm.html
188 : :
189 : 0 : tk::real B[m_ncomp][m_ncomp];
190 : 0 : tk::real Bo[m_ncomp][m_ncomp];
191 [ - - ]: 0 : for (ncomp_t i=0; i<m_ncomp; ++i) {
192 : 0 : const tk::real& pari = particles( p, i, m_offset );
193 [ - - ]: 0 : for (ncomp_t j=0; j<m_ncomp; ++j) {
194 : : const tk::real& parj = particles( p, j, m_offset );
195 [ - - ]: 0 : if (i == j) {
196 [ - - ]: 0 : B[i][i] = std::abs( pari * (1.0 - pari) );
197 [ - - ]: 0 : if (B[i][i] < 1.0e-10) B[i][i] = 1.0;
198 : : } else {
199 : 0 : B[i][j] = -pari*parj;
200 : : }
201 : : }
202 : : }
203 : 0 : std::memcpy( Bo, B, m_ncomp*m_ncomp*sizeof(tk::real) );
204 : : // Compute diffusion matrix (lower triangle of Cholesky-decomposition)
205 : : lapack_int n = static_cast< lapack_int >( m_ncomp );
206 : 0 : lapack_int info = LAPACKE_dpotrf( LAPACK_ROW_MAJOR, 'L', n, B[0], n );
207 [ - - ]: 0 : if (info != 0 ) {
208 : 0 : print_matrix( "=======\nOriginal Matrix", Bo[0], n, n );
209 : : std::cout << "info: " << info << std::endl;
210 : 0 : print_matrix( "Result of Cholesky factorization", B[0], n, n );
211 [ - - ]: 0 : for (ncomp_t i=0; i<m_ncomp; ++i) {
212 : : std::cout <<
213 : : std::setprecision( std::numeric_limits< tk::real >::digits10 )
214 : 0 : << i << " par: " << particles( p, i, m_offset) << std::endl;
215 : : }
216 : : }
217 : : //ErrChk( info == 0, "Wright-Fisher Cholesky decomposition unsuccessful, "
218 : : // "info = " + std::to_string( info ) + ", particle: " +
219 : : // std::to_string( p ) );
220 : :
221 : : // Advance the first m_ncomp (N-1) scalars
222 [ - - ]: 0 : if (info == 0) {
223 : : ncomp_t i = 0;
224 [ - - ]: 0 : for (i=0; i<m_ncomp-1; ++i) {
225 : 0 : tk::real& par = particles( p, i, m_offset );
226 : : // Advance first m_ncomp (K=N-1) scalars due to drift
227 : 0 : par += 0.5*(m_omega[i] - omega*par)*dt;
228 : : // Advance first m_ncomp (K=N-1) particles with Cholesky-decomposed
229 : : // lower triangle (diffusion matrix)
230 [ - - ]: 0 : for (ncomp_t j=0; j<m_ncomp-1; ++j)
231 [ - - ]: 0 : if (j<=i) {
232 : : tk::real dW;
233 : 0 : m_rng.gaussian( stream, 1, &dW );
234 : 0 : par += B[i][j] * sqrt(dt) * dW;
235 : : }
236 : : }
237 : : // Compute the (N-1)th scalar from unit-sum
238 : 0 : tk::real& par = particles( p, i, m_offset );
239 : 0 : par = 1.0 - particles( p, 0, m_offset );
240 [ - - ]: 0 : for (i=1; i<m_ncomp-1; ++i) {
241 : 0 : par -= particles( p, i, m_offset );
242 : : }
243 : : }
244 : : }
245 : :
246 : : #if defined(__clang__)
247 : : #pragma clang diagnostic pop
248 : : #elif defined(STRICT_GNUC)
249 : : #pragma GCC diagnostic pop
250 : : #endif
251 : 0 : }
252 : :
253 : : private:
254 : : const ncomp_t m_c; //!< Equation system index
255 : : const char m_depvar; //!< Dependent variable
256 : : const ncomp_t m_ncomp; //!< Number of components
257 : : const ncomp_t m_offset; //!< Offset SDE operates from
258 : : const tk::RNG& m_rng; //!< Random number generator
259 : :
260 : : //! Coefficients
261 : : std::vector< kw::sde_omega::info::expect::type > m_omega;
262 : :
263 : : //! Coefficients policy
264 : : Coefficients coeff;
265 : : };
266 : :
267 : : } // walker::
268 : :
269 : : #endif // WrightFisher_h
|
