Commit remaining bits
[pspp] / src / language / xforms / sample.c
1 /* PSPP - a program for statistical analysis.
2    Copyright (C) 1997-9, 2000, 2009-2011 Free Software Foundation, Inc.
3
4    This program is free software: you can redistribute it and/or modify
5    it under the terms of the GNU General Public License as published by
6    the Free Software Foundation, either version 3 of the License, or
7    (at your option) any later version.
8
9    This program is distributed in the hope that it will be useful,
10    but WITHOUT ANY WARRANTY; without even the implied warranty of
11    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12    GNU General Public License for more details.
13
14    You should have received a copy of the GNU General Public License
15    along with this program.  If not, see <http://www.gnu.org/licenses/>. */
16
17 #include <config.h>
18
19 #include <gsl/gsl_rng.h>
20 #include <limits.h>
21 #include <stdio.h>
22 #include <math.h>
23
24 #include "data/dataset.h"
25 #include "data/variable.h"
26 #include "language/command.h"
27 #include "language/lexer/lexer.h"
28 #include "libpspp/compiler.h"
29 #include "libpspp/message.h"
30 #include "libpspp/str.h"
31 #include "math/random.h"
32
33 #include "gl/xalloc.h"
34
35 #include "gettext.h"
36 #define _(msgid) gettext (msgid)
37
38 /* The two different types of samples. */
39 enum
40   {
41     TYPE_A_FROM_B,              /* 5 FROM 10 */
42     TYPE_FRACTION               /* 0.5 */
43   };
44
45 /* SAMPLE transformation. */
46 struct sample_trns
47   {
48     int type;                   /* One of TYPE_*. */
49     int n, N;                   /* TYPE_A_FROM_B: n from N. */
50     int m, t;                   /* TYPE_A_FROM_B: # picked so far; # so far. */
51     unsigned frac;              /* TYPE_FRACTION: a fraction of UINT_MAX. */
52   };
53
54 static trns_proc_func sample_trns_proc;
55 static trns_free_func sample_trns_free;
56
57 int
58 cmd_sample (struct lexer *lexer, struct dataset *ds)
59 {
60   struct sample_trns *trns;
61
62   int type;
63   int a, b;
64   unsigned frac;
65
66   if (!lex_force_num (lexer))
67     return CMD_FAILURE;
68   if (!lex_is_integer (lexer))
69     {
70       unsigned long min = gsl_rng_min (get_rng ());
71       unsigned long max = gsl_rng_max (get_rng ());
72
73       type = TYPE_FRACTION;
74       if (lex_tokval (lexer) <= 0 || lex_tokval (lexer) >= 1)
75         {
76           msg (SE, _("The sampling factor must be between 0 and 1 "
77                      "exclusive."));
78           return CMD_FAILURE;
79         }
80
81       frac = lex_tokval (lexer) * (max - min) + min;
82       a = b = 0;
83     }
84   else
85     {
86       type = TYPE_A_FROM_B;
87       a = lex_integer (lexer);
88       lex_get (lexer);
89       if (!lex_force_match_id (lexer, "FROM"))
90         return CMD_FAILURE;
91       if (!lex_force_int (lexer))
92         return CMD_FAILURE;
93       b = lex_integer (lexer);
94       if (a >= b)
95         {
96           msg (SE, _("Cannot sample %d observations from a population of "
97                      "%d."),
98                a, b);
99           return CMD_FAILURE;
100         }
101
102       frac = 0;
103     }
104   lex_get (lexer);
105
106   trns = xmalloc (sizeof *trns);
107   trns->type = type;
108   trns->n = a;
109   trns->N = b;
110   trns->m = trns->t = 0;
111   trns->frac = frac;
112   add_transformation (ds, sample_trns_proc, sample_trns_free, trns);
113
114   return CMD_SUCCESS;
115 }
116
117 /* Executes a SAMPLE transformation. */
118 static int
119 sample_trns_proc (void *t_, struct ccase **c UNUSED,
120                   casenumber case_num UNUSED)
121 {
122   struct sample_trns *t = t_;
123   double U;
124
125   if (t->type == TYPE_FRACTION)
126     {
127       if (gsl_rng_get (get_rng ()) <= t->frac)
128         return TRNS_CONTINUE;
129       else
130         return TRNS_DROP_CASE;
131     }
132
133   if (t->m >= t->n)
134     return TRNS_DROP_CASE;
135
136   U = gsl_rng_uniform (get_rng ());
137   if ((t->N - t->t) * U >= t->n - t->m)
138     {
139       t->t++;
140       return TRNS_DROP_CASE;
141     }
142   else
143     {
144       t->m++;
145       t->t++;
146       return TRNS_CONTINUE;
147     }
148 }
149
150 static bool
151 sample_trns_free (void *t_)
152 {
153   struct sample_trns *t = t_;
154   free (t);
155   return true;
156 }