2 files changed, 52 insertions, 0 deletions

diff --git a/gi/pf/align-lexonly-pyp.cc b/gi/pf/align-lexonly-pyp.cc
index 4a1d1db6..0c90b6ce 100644
--- a/gi/pf/align-lexonly-pyp.cc
+++ b/gi/pf/align-lexonly-pyp.cc
@@ -11,6 +11,7 @@
 #include "sampler.h"
 #include "corpus.h"
 #include "pyp_tm.h"
+#include "quasi_model2.h"
 
 using namespace std;
 namespace po = boost::program_options;
@@ -61,12 +62,14 @@ struct Aligner {
   Aligner(const vector<vector<WordID> >& lets, int num_letters, vector<AlignedSentencePair>* c) :
       corpus(*c),
       model(lets, num_letters),
+      paj(4, 0.08),
       kNULL(TD::Convert("NULL")) {
     assert(lets[kNULL].size() == 0);
   }
 
   vector<AlignedSentencePair>& corpus;
   PYPLexicalTranslation model;
+  const QuasiModel2 paj;
   const WordID kNULL;
 
   void ResampleHyperparameters() {
@@ -83,6 +86,7 @@ struct Aligner {
         a_j = prng->next() * (1 + asp.src.size());
         const WordID f_a_j = (a_j ? asp.src[a_j - 1] : kNULL);
         model.Increment(f_a_j, asp.trg[j], &*prng);
+        // TODO factor in alignment prob
       }
     }
     cerr << "Corpus intialized randomly. LLH = " << model.Likelihood() << endl;
@@ -101,6 +105,8 @@ struct Aligner {
         for (unsigned prop_a_j = 0; prop_a_j <= asp.src.size(); ++prop_a_j) {
           const WordID prop_f = (prop_a_j ? asp.src[prop_a_j - 1] : kNULL);
           ss[prop_a_j] = model.Prob(prop_f, e_j);
+          // TODO configurable
+          ss[prop_a_j] *= paj.Pa_j(prop_a_j, j, asp.src.size(), asp.trg.size());
         }
         a_j = prng->SelectSample(ss);
         f_a_j = (a_j ? asp.src[a_j - 1] : kNULL);
diff --git a/gi/pf/quasi_model2.h b/gi/pf/quasi_model2.h
new file mode 100644
index 00000000..0095289f
--- /dev/null
+++ b/gi/pf/quasi_model2.h
@@ -0,0 +1,46 @@
+#ifndef _QUASI_MODEL2_H_
+#define _QUASI_MODEL2_H_
+
+#include <vector>
+#include <cmath>
+#include "prob.h"
+#include "array2d.h"
+
+struct QuasiModel2 {
+  explicit QuasiModel2(double alpha, double pnull = 0.1) :
+      alpha_(alpha),
+      pnull_(pnull),
+      pnotnull_(1 - pnull),
+      z_(1000,1000) {}
+  // a_j = 0 => NULL; src_len does *not* include null
+  prob_t Pa_j(unsigned a_j, unsigned j, unsigned src_len, unsigned trg_len) const {
+    if (!a_j) return pnull_;
+    std::vector<prob_t>& zv = z_(src_len, trg_len);
+    if (zv.size() == 0)
+      zv.resize(trg_len);
+    
+    prob_t& z = zv[j];
+    if (z.is_0()) z = ComputeZ(j, src_len, trg_len);
+
+    prob_t p;
+    p.logeq(-fabs(double(a_j - 1) / src_len - double(j) / trg_len) * alpha_);
+    p *= pnotnull_;
+    p /= z;
+    return p;
+  }
+ private:
+  prob_t ComputeZ(unsigned j, unsigned src_len, unsigned trg_len) const {
+    prob_t p, z = prob_t::Zero();
+    for (int a_j = 1; a_j <= src_len; ++a_j) {
+      p.logeq(-fabs(double(a_j - 1) / src_len - double(j) / trg_len) * alpha_);
+      z += p;
+    }
+    return z;
+  }
+  double alpha_;
+  const prob_t pnull_;
+  const prob_t pnotnull_;
+  mutable Array2D<std::vector<prob_t> > z_;
+};
+
+#endif